"""Cost report — appends each run to a JSON sidecar and regenerates cost.md."""

from __future__ import annotations

import json
import logging
from datetime import datetime, timezone
from pathlib import Path

from secagent.models import TokenUsage, _MODEL_PRICING, _DEFAULT_PRICING

logger = logging.getLogger(__name__)

_STAGE_ORDER = ["planner", "scanner", "verifier", "fixer", "validator"]


def update_cost_report(
    stage_usage: dict[str, TokenUsage],
    stage_models: dict[str, str],
    command: str,
    target: str,
    output_dir: Path,
) -> Path:
    """Append this run to the cost sidecar and regenerate ``cost.md``.

    Parameters
    ----------
    stage_usage:
        Per-stage token usage from the pipeline run.
    stage_models:
        Per-stage model name (e.g. planner/verifier use the fast model,
        scanner/fixer/validator use the reasoning model).
    command:
        ``"scan"`` or ``"fix"``.
    target:
        Path to the scanned/fixed target.
    output_dir:
        The ``.secagent/`` cache directory.

    Returns
    -------
    Path
        Absolute path to the updated ``cost.md``.
    """
    data_dir = output_dir / "data"
    data_dir.mkdir(parents=True, exist_ok=True)
    data_file = data_dir / "cost_data.json"
    runs: list[dict] = []
    if data_file.exists():
        try:
            runs = json.loads(data_file.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            runs = []

    # Build the new run record — model stored per stage for accurate pricing
    run: dict = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "command": command,
        "target": target,
        "stages": {
            stage: {
                "model": stage_models.get(stage, ""),
                "prompt_tokens": u.prompt_tokens,
                "completion_tokens": u.completion_tokens,
                "reasoning_tokens": u.reasoning_tokens,
                "calls": u.calls,
            }
            for stage, u in stage_usage.items()
        },
    }
    runs.append(run)
    data_file.write_text(json.dumps(runs, indent=2), encoding="utf-8")

    report_path = output_dir / "cost.md"
    _write_cost_md(runs, report_path)
    logger.info("Cost report updated: %s", report_path)
    return report_path


# ── Markdown generation ───────────────────────────────────────────────────────

def _stage_cost(sd: dict) -> float:
    """Compute cost for a single stage record using its per-stage model."""
    model = sd.get("model", "")
    in_price, out_price = _MODEL_PRICING.get(model, _DEFAULT_PRICING)
    p = sd.get("prompt_tokens", 0)
    c = sd.get("completion_tokens", 0)
    return (p * in_price + c * out_price) / 1_000_000


def _write_cost_md(runs: list[dict], report_path: Path) -> None:
    """Regenerate the full cost.md from the list of run records."""
    has_reasoning = any(
        sd.get("reasoning_tokens", 0) > 0
        for run in runs
        for sd in run.get("stages", {}).values()
    )

    # Collect unique models used across all runs for the header
    all_models: set[str] = set()
    for run in runs:
        for sd in run["stages"].values():
            m = sd.get("model", "")
            if m:
                all_models.add(m)

    lines: list[str] = [
        "# SecAgent API Cost Log\n",
        f"**Target:** `{runs[0]['target']}`  ",
        f"**Models:** {', '.join(f'`{m}`' for m in sorted(all_models))}\n",
    ]

    if has_reasoning:
        lines.append("> ⚡ Reasoning tokens are a subset of output tokens, billed at the output rate.\n")

    # ── Cumulative totals ─────────────────────────────────────────────
    cum_prompt = cum_completion = cum_reasoning = cum_calls = 0
    cum_cost = 0.0
    for run in runs:
        for sd in run["stages"].values():
            cum_prompt += sd.get("prompt_tokens", 0)
            cum_completion += sd.get("completion_tokens", 0)
            cum_reasoning += sd.get("reasoning_tokens", 0)
            cum_calls += sd.get("calls", 0)
            cum_cost += _stage_cost(sd)

    if has_reasoning:
        lines += [
            "## Cumulative Totals\n",
            "| Input | Output | Reasoning *(in output)* | Total | Calls | Est. Cost |",
            "|---|---|---|---|---|---|",
            f"| {cum_prompt:,} | {cum_completion:,} | {cum_reasoning:,} | {cum_prompt + cum_completion:,} | {cum_calls} | **${cum_cost:.4f}** |",
        ]
    else:
        lines += [
            "## Cumulative Totals\n",
            "| Input | Output | Total | Calls | Est. Cost |",
            "|---|---|---|---|---|",
            f"| {cum_prompt:,} | {cum_completion:,} | {cum_prompt + cum_completion:,} | {cum_calls} | **${cum_cost:.4f}** |",
        ]

    lines += ["", "---\n", "## Run History\n"]

    # ── Per-run breakdown (oldest → newest) ──────────────────────────
    for idx, run in enumerate(runs, start=1):
        ts = run.get("timestamp", "unknown")
        try:
            dt = datetime.fromisoformat(ts).strftime("%Y-%m-%d %H:%M UTC")
        except ValueError:
            dt = ts

        run_prompt = run_completion = run_reasoning = run_calls = 0
        run_cost = 0.0
        for sd in run["stages"].values():
            run_prompt += sd.get("prompt_tokens", 0)
            run_completion += sd.get("completion_tokens", 0)
            run_reasoning += sd.get("reasoning_tokens", 0)
            run_calls += sd.get("calls", 0)
            run_cost += _stage_cost(sd)

        lines += [f"### Run {idx} — `{run.get('command', '?')}` — {dt}\n"]

        if has_reasoning:
            lines += [
                "| Stage | Model | Input | Output | Reasoning | Total | Calls | Cost |",
                "|---|---|---|---|---|---|---|---|",
            ]
        else:
            lines += [
                "| Stage | Model | Input | Output | Total | Calls | Cost |",
                "|---|---|---|---|---|---|---|",
            ]

        stages = run.get("stages", {})
        for stage in _STAGE_ORDER:
            if stage not in stages:
                continue
            sd = stages[stage]
            p = sd.get("prompt_tokens", 0)
            c = sd.get("completion_tokens", 0)
            r = sd.get("reasoning_tokens", 0)
            calls = sd.get("calls", 0)
            model = sd.get("model", "—")
            cost = _stage_cost(sd)
            if has_reasoning:
                lines.append(
                    f"| {stage} | `{model}` | {p:,} | {c:,} | {r:,} | {p + c:,} | {calls} | ${cost:.4f} |"
                )
            else:
                lines.append(
                    f"| {stage} | `{model}` | {p:,} | {c:,} | {p + c:,} | {calls} | ${cost:.4f} |"
                )

        if has_reasoning:
            lines.append(
                f"| **Total** | | **{run_prompt:,}** | **{run_completion:,}** | **{run_reasoning:,}** | **{run_prompt + run_completion:,}** | **{run_calls}** | **${run_cost:.4f}** |"
            )
        else:
            lines.append(
                f"| **Total** | | **{run_prompt:,}** | **{run_completion:,}** | **{run_prompt + run_completion:,}** | **{run_calls}** | **${run_cost:.4f}** |"
            )
        lines.append("")

    report_path.write_text("\n".join(lines), encoding="utf-8")
