#!/usr/bin/env python3
"""pragent pilot — per-review cost model.

Answers "what would this cost on a paid API?" for the pilot's agent loop. The
pilot currently runs on `glm-5.2:cloud` through the on-network headroom proxy at
no per-token charge, so every review's measured usage is *free but real*: it
tells us exactly what the same work would bill on Claude or GPT.

The model is deliberately explicit rather than a single fudge factor, because
the dominant cost in an agent loop is not the diff — it is **resending the
conversation on every step**. A 12-step review re-reads its own prefix 12 times.
Prompt caching is what makes that affordable, and whether caching is on changes
the answer by ~3x, so it's a parameter, not an assumption.

Token accounting per review:

    step 1 input  = prefix + brief
    step k input  = prefix + brief + (tool results accumulated through k-1)
    total input   = sum over steps
    cached        = the prefix + brief part of steps 2..n (stable, byte-identical)
    uncached      = step 1 in full + the growing tool-result tail

`prefix` = system + tool schemas + agent definition + the skills this tier loads.
Those sizes are MEASURED from the files in this repo (see `measure_factory`),
not guessed. Diff size, file reads, and step count are per-tier assumptions from
the `attention-tiering` skill's budgets — override them on the CLI to fit your
own repos.

Prices are per million tokens, from the providers' published pricing pages
(fetched 2026-08-18 — re-check before quoting):
  https://platform.claude.com/docs/en/about-claude/pricing
  https://developers.openai.com/api/docs/pricing

Usage:
    python3 pilot/cost_model.py                       # all tiers, all models
    python3 pilot/cost_model.py --prs-per-month 350
    python3 pilot/cost_model.py --mix 5,35,55,5       # trivial,lite,full,oversized %
    python3 pilot/cost_model.py --no-cache            # what caching is worth
"""

from __future__ import annotations

import argparse
import os
from dataclasses import dataclass, field

CHARS_PER_TOKEN = 4  # English prose/code rule of thumb; ±15% is normal


# ---------------------------------------------------------------------------
# Prices — USD per million tokens
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class Price:
    """Per-MTok prices. `cache_write` and `cache_read` are absolute rates, not
    multipliers, so providers with different cache economics stay comparable."""

    name: str
    input: float
    output: float
    cache_write: float
    cache_read: float

    @property
    def batch_input(self) -> float:
        return self.input / 2

    @property
    def batch_output(self) -> float:
        return self.output / 2


# Anthropic: cache write = 1.25x input (5-minute TTL), cache read = 0.1x input.
# OpenAI: cached input is a published rate (0.1x input); there is no separate
# cache-write charge — writes are billed as ordinary input.
PRICES: dict[str, Price] = {
    "claude-opus-5": Price("Claude Opus 5", 5.00, 25.00, 6.25, 0.50),
    "claude-sonnet-5": Price("Claude Sonnet 5", 2.00, 10.00, 2.50, 0.20),
    "claude-haiku-4-5": Price("Claude Haiku 4.5", 1.00, 5.00, 1.25, 0.10),
    "gpt-5.6-sol": Price("GPT-5.6 Sol", 5.00, 30.00, 5.00, 0.50),
    "gpt-5.6-terra": Price("GPT-5.6 Terra", 2.00, 12.00, 2.00, 0.20),
    "gpt-5.6-luna": Price("GPT-5.6 Luna", 0.20, 1.20, 0.20, 0.02),
}


# ---------------------------------------------------------------------------
# Factory footprint — measured from this repo
# ---------------------------------------------------------------------------

_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# Skills the primary always loads, and the conditional ones per tier. Mirrors
# the load table in .opencode/agents/pragent.md.
ALWAYS_SKILLS = ("review-methodology", "findings-schema", "attention-tiering")
TIER_SKILLS: dict[str, tuple[str, ...]] = {
    "trivial": (),
    "lite": ("comment-craft",),
    "full": ("linter-playbook", "security-lens", "comment-craft"),
    "oversized": ("linter-playbook", "security-lens", "comment-craft", "malicious-change"),
}

# opencode's own system prompt + the JSON tool schemas it sends (read, grep,
# glob, bash, webfetch, skill, task, …). Not in this repo, so this is the one
# component that is an estimate rather than a measurement.
HARNESS_TOKENS = 3500


def _tok(path: str) -> int:
    try:
        with open(path, "rb") as f:
            return len(f.read()) // CHARS_PER_TOKEN
    except OSError:
        return 0


def measure_factory(root: str = _ROOT) -> dict[str, int]:
    """Token size of each prompt component, measured from the files on disk."""
    out = {"agent": _tok(os.path.join(root, ".opencode", "agents", "pragent.md"))}
    skills_dir = os.path.join(root, ".opencode", "skills")
    if os.path.isdir(skills_dir):
        for name in sorted(os.listdir(skills_dir)):
            p = os.path.join(skills_dir, name, "SKILL.md")
            if os.path.isfile(p):
                out[f"skill:{name}"] = _tok(p)
    for lens in ("security", "tests", "perf"):
        out[f"subagent:{lens}"] = _tok(os.path.join(root, ".opencode", "agents", f"{lens}.md"))
    return out


def prefix_tokens(tier: str, factory: dict[str, int]) -> int:
    """Stable per-step prefix: harness + agent definition + loaded skills."""
    total = HARNESS_TOKENS + factory.get("agent", 0)
    for s in ALWAYS_SKILLS + TIER_SKILLS.get(tier, ()):
        total += factory.get(f"skill:{s}", 0)
    return total


# ---------------------------------------------------------------------------
# Per-tier workload assumptions
# ---------------------------------------------------------------------------


@dataclass
class Tier:
    """One tier's workload. Defaults follow the `attention-tiering` budgets."""

    name: str
    diff_tokens: int          # the diff as it lands in the brief
    steps: int                # model turns in the agent loop
    file_reads: int           # files read from the checkout
    tokens_per_read: int      # avg tokens returned per read/grep/linter result
    output_tokens: int        # assistant output across all steps (incl. reasoning)
    subagents: int = 0        # lens subagents spawned
    brief_fixed: int = 600    # brief template + PR meta + prior reviews
    share: float = 0.0        # fraction of PRs at this tier (for the monthly mix)
    _factory: dict = field(default_factory=dict, repr=False)


DEFAULT_TIERS = [
    # diff_tok  steps  reads  tok/read  output  subs  share
    Tier("trivial",     400,  2,  0,     0,      600,  0, share=0.05),
    Tier("lite",       1500,  6,  4,  2000,     2500,  0, share=0.35),
    Tier("full",       6000, 24, 20,  3300,    12000,  0, share=0.55),
    Tier("oversized", 25000, 35, 30,  3500,    20000,  2, share=0.05),
]


# ---------------------------------------------------------------------------
# Observed runs — the calibration anchor
# ---------------------------------------------------------------------------

# Real usage reported by the AI-USAGE label, summed from opencode's step_finish
# events. Keep this list append-only: it is the only thing separating this model
# from a guess, and the first entry corrected the tier assumptions by ~15x.
OBSERVED_RUNS: list[dict] = [
    {
        "label": "internal/hardening-PR (16 files, 1020 insertions / 91 deletions)",
        "date": "2026-08-18",
        "tier": "full",
        "diff_tokens": 17_600,     # 16 files, 1020 insertions / 91 deletions
        "steps": 28,
        "duration_s": 348.3,
        "input": 2_071_025,
        "output": 17_303,
        "cache_read": 0,
        "cache_write": 0,
        "subagents": 0,
    },
    {
        "label": "internal/hardening-PR (same PR, two commits later)",
        "date": "2026-08-18",
        "tier": "full",
        "diff_tokens": 21_000,     # same PR, two commits later
        "steps": 31,
        "duration_s": 189.8,
        "input": 2_213_077,
        "output": 9_058,
        "cache_read": 0,
        "cache_write": 0,
        "subagents": 0,
    },
    # A third run of the same PR (sha 2613b3e, 31 steps' worth of work in 330s)
    # ended without a parseable findings block and so reported no usage at all —
    # the reason `salvage_summary` now keeps the usage section on that path.
]


def observed_usage(run: dict) -> Usage:
    return Usage(
        uncached_input=run["input"] - run.get("cache_read", 0),
        cached_input=run.get("cache_read", 0),
        cache_writes=run.get("cache_write", 0),
        output=run["output"],
    )


@dataclass
class Usage:
    uncached_input: int = 0
    cached_input: int = 0
    cache_writes: int = 0
    output: int = 0

    @property
    def total_input(self) -> int:
        return self.uncached_input + self.cached_input


def tier_usage(tier: Tier, factory: dict[str, int], caching: bool = True) -> Usage:
    """Token usage for one review at this tier.

    The agent loop resends the whole conversation each step. The prefix + brief
    are byte-identical across steps, so with caching they are written once and
    read back on every later step; the tool-result tail grows and is charged as
    ordinary input. Without caching every step pays full input price for
    everything it has accumulated — which is the quadratic term that makes an
    uncached agent loop expensive.
    """
    prefix = prefix_tokens(tier.name, factory)
    stable = prefix + tier.brief_fixed + tier.diff_tokens

    # Tool results arrive one per step, after the first.
    result_steps = max(0, min(tier.file_reads, tier.steps - 1))
    per_result = tier.tokens_per_read

    u = Usage(output=tier.output_tokens)

    if caching:
        u.cache_writes = stable
        u.cached_input = stable * max(0, tier.steps - 1)
        u.uncached_input = 0
    else:
        u.uncached_input = stable * tier.steps

    # The growing tail of tool results: a result produced at step i is resent on
    # every step after it, so it is counted (steps - i) times.
    tail = 0
    for i in range(1, result_steps + 1):
        tail += per_result * (tier.steps - i)
    u.uncached_input += tail

    # Each lens subagent is its own loop: its own prefix, the diff, a few reads.
    for _ in range(tier.subagents):
        sub_prefix = HARNESS_TOKENS + factory.get("subagent:security", 600)
        sub_stable = sub_prefix + tier.diff_tokens
        sub_steps = 6
        if caching:
            u.cache_writes += sub_stable
            u.cached_input += sub_stable * (sub_steps - 1)
        else:
            u.uncached_input += sub_stable * sub_steps
        for i in range(1, 4):
            u.uncached_input += per_result * (sub_steps - i)
        u.output += 1500

    return u


def cost(u: Usage, price: Price, batch: bool = False) -> float:
    """USD for one review's usage at these prices."""
    inp = price.batch_input if batch else price.input
    out = price.batch_output if batch else price.output
    cw = price.cache_write / 2 if batch else price.cache_write
    cr = price.cache_read / 2 if batch else price.cache_read
    return (
        u.uncached_input * inp
        + u.cached_input * cr
        + u.cache_writes * cw
        + u.output * out
    ) / 1_000_000


# ---------------------------------------------------------------------------
# Reporting
# ---------------------------------------------------------------------------


def blended_cost(tiers: list[Tier], factory: dict, price: Price, caching: bool) -> float:
    """Weighted cost of one average PR across the tier mix."""
    total_share = sum(t.share for t in tiers) or 1.0
    return sum(
        cost(tier_usage(t, factory, caching), price) * (t.share / total_share)
        for t in tiers
    )


def report(tiers: list[Tier], prs_per_month: int, caching: bool, models: list[str]) -> str:
    factory = measure_factory()
    lines: list[str] = []

    lines.append(f"Factory footprint (measured, {CHARS_PER_TOKEN} chars/token):")
    for k, v in sorted(factory.items()):
        lines.append(f"  {k:<34} {v:>6,} tok")
    lines.append(f"  {'harness (opencode + tool schemas, est.)':<34} {HARNESS_TOKENS:>6,} tok")
    lines.append("")

    lines.append(f"Per-review tokens  (prompt caching: {'on' if caching else 'OFF'})")
    lines.append(f"  {'tier':<11} {'prefix':>8} {'uncached':>10} {'cached':>10} {'cwrite':>8} {'output':>8}")
    for t in tiers:
        u = tier_usage(t, factory, caching)
        lines.append(
            f"  {t.name:<11} {prefix_tokens(t.name, factory):>8,} {u.uncached_input:>10,} "
            f"{u.cached_input:>10,} {u.cache_writes:>8,} {u.output:>8,}"
        )
    lines.append("")

    lines.append("Cost per review (USD)")
    header = f"  {'model':<18}" + "".join(f"{t.name:>12}" for t in tiers) + f"{'blended':>12}"
    lines.append(header)
    for key in models:
        p = PRICES[key]
        row = f"  {p.name:<18}"
        for t in tiers:
            row += f"{cost(tier_usage(t, factory, caching), p):>12.4f}"
        row += f"{blended_cost(tiers, factory, p, caching):>12.4f}"
        lines.append(row)
    lines.append("")

    mix = ", ".join(f"{t.name} {t.share:.0%}" for t in tiers)
    lines.append(f"Monthly at {prs_per_month} PRs/month  (mix: {mix})")
    lines.append(f"  {'model':<18} {'per PR':>10} {'per month':>12} {'batch -50%':>12}")
    for key in models:
        p = PRICES[key]
        per_pr = blended_cost(tiers, factory, p, caching)
        lines.append(
            f"  {p.name:<18} {per_pr:>10.4f} {per_pr * prs_per_month:>12.2f}"
            f" {per_pr * prs_per_month / 2:>12.2f}"
        )
    lines.append("")
    lines.append("Batch column applies the 50% async discount; it is shown for scale only —")
    lines.append("PR review is latency-sensitive and a stateful agent loop is not batchable.")
    lines.append("")
    lines.append(observed_report(models))
    return "\n".join(lines)


def observed_report(models: list[str]) -> str:
    """Price the runs actually measured through the AI-USAGE label."""
    if not OBSERVED_RUNS:
        return "No observed runs recorded yet."
    lines = ["Observed runs (measured via the AI-USAGE label)"]
    for run in OBSERVED_RUNS:
        u = observed_usage(run)
        lines.append(
            f"  {run['label']} — tier {run['tier']}, {run['steps']} steps, "
            f"{run['duration_s']:.0f}s, {run['input']:,} in / {run['output']:,} out, "
            f"cache {run['cache_read']:,} read / {run['cache_write']:,} write"
        )
        row = "   "
        for key in models:
            p = PRICES[key]
            row += f" {p.name}: ${cost(u, p):.2f} "
        lines.append(row)
    lines.append("")
    lines.append("  NOTE: the pilot's headroom/glm-5.2 path reports zero cache read and zero")
    lines.append("  cache write, i.e. prompt caching is NOT in play today. On a provider where")
    lines.append("  it is, the stable prefix (agent + skills + brief + diff, resent every step)")
    lines.append("  drops to 0.1x — worth roughly a third of the bill on a run like the one")
    lines.append("  above. Budget with caching OFF until the measured cache columns are nonzero.")
    return "\n".join(lines)


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="pragent per-review cost model")
    ap.add_argument("--prs-per-month", type=int, default=350)
    ap.add_argument("--mix", default="", help="trivial,lite,full,oversized as percentages")
    ap.add_argument("--no-cache", action="store_true", help="model without prompt caching")
    ap.add_argument("--models", default=",".join(PRICES))
    args = ap.parse_args(argv)

    tiers = DEFAULT_TIERS
    if args.mix:
        shares = [float(x) for x in args.mix.split(",")]
        if len(shares) != len(tiers):
            ap.error(f"--mix needs {len(tiers)} comma-separated values")
        for t, s in zip(tiers, shares):
            t.share = s / 100.0

    models = [m.strip() for m in args.models.split(",") if m.strip()]
    unknown = [m for m in models if m not in PRICES]
    if unknown:
        ap.error(f"unknown model(s): {', '.join(unknown)}")

    print(report(tiers, args.prs_per_month, not args.no_cache, models))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())