"""Final-report data.json mutation helpers (Phase 7 token cells)."""
from __future__ import annotations

import json
import sys
from pathlib import Path
from typing import Any

# Make sibling packages importable when this module runs as part of an
# in-repo invocation. The installed runtime puts everything under
# ~/.okstra/lib/python which is already on PYTHONPATH via the bin
# wrapper, so this path tweak is a no-op there.
_HERE = Path(__file__).resolve().parent.parent
if str(_HERE) not in sys.path:
    sys.path.insert(0, str(_HERE))

from okstra_ctl.design_prep import (  # noqa: E402
    DesignPrepError,
    materialize_design_prep_requests,
)
from okstra_ctl.usage_cells import duration_ms_from_bounds  # noqa: E402

from .task_totals import task_cumulative_usage  # noqa: E402


class SubstituteRefusedError(RuntimeError):
    """Raised when substitution would write zero-only token cells.

    Shipping `0` / `$0.00` in the Lead / Worker / Grand rows is the
    observed silent-failure mode where the collector ran but every
    session jsonl was empty (or the writer fabricated zeros). Surfacing
    the refusal here lets the caller retry after locating the missing
    session jsonls instead of baking zeros into the report.
    """


def _norm(s: str) -> str:
    """alnum-only lowercase, so role / agent / model line up across
    rephrasings (``Claude worker`` vs ``claude``; ``gpt-5.6-sol`` vs
    ``gpt56sol``)."""
    return "".join(ch.lower() for ch in s if ch.isalnum())


def _match_worker_index(row: dict, workers: list, used: set[int]) -> int | None:
    """Pick the team-state worker that owns this executionStatus row, never
    reusing a worker an earlier row already claimed.

    A contract-following report names each row after its team-state worker
    (``Claude worker``), so role equality / containment matches. A report
    that instead used function roles (``Analysis verifier``, ``Acceptance
    critic``) falls back to a *unique* provider(agent)+model match. When two
    rows share one provider aggregate (codex verifier + critic → one ``Codex
    worker``), only the first is attributed; the second stays null because
    team-state holds no separate figure for it.
    """
    role = _norm(row.get("role", ""))
    agent = _norm(row.get("agent", ""))
    model = _norm(row.get("model", ""))
    for i, worker in enumerate(workers):
        if i in used:
            continue
        worker_role = _norm(worker.get("role", ""))
        if worker_role and role and (
            worker_role == role or worker_role in role or role in worker_role
        ):
            return i
    candidates = [
        i
        for i, worker in enumerate(workers)
        if i not in used
        and _norm(worker.get("agent", ""))
        and _norm(worker.get("agent", "")) in agent
        and _norm(worker.get("model", "")) == model
    ]
    return candidates[0] if len(candidates) == 1 else None


def _populate_execution_row(row: dict, source: dict) -> None:
    """Fill the executionStatus row's token / cost / duration cells from a
    team-state worker (or lead) entry. ``source`` is the dict that owns
    a ``usage`` sub-dict and (for workers) a ``status`` field.
    """
    usage = source.get("usage") or {}
    if usage.get("source") == "unavailable":
        # Token cells stay null (`--`). Dispatch wall-clock still belongs
        # on the row: a wrapper that failed still ran for some time.
        duration = duration_ms_from_bounds(source.get("startedAt"), source.get("endedAt"))
        if duration is not None:
            row["durationMs"] = duration
        return
    if "totalTokens" in usage:
        row["totalTokens"] = usage["totalTokens"]
    if "cacheReadTokens" in usage:
        row["cacheReadTokens"] = usage["cacheReadTokens"]
    if "billableEquivalentTokens" in usage:
        row["billableTokens"] = usage["billableEquivalentTokens"]
    if "estimatedCostUsd" in usage:
        row["costUsd"] = usage["estimatedCostUsd"]
    if "durationMs" in usage:
        row["durationMs"] = usage["durationMs"]
    if "cliTotalTokens" in usage:
        row["cliTotalTokens"] = usage["cliTotalTokens"]
    if "cliEstimatedCostUsd" in usage:
        row["cliCostUsd"] = usage["cliEstimatedCostUsd"]
    if not row.get("durationMs"):
        duration = duration_ms_from_bounds(source.get("startedAt"), source.get("endedAt"))
        if duration is None:
            duration = duration_ms_from_bounds(usage.get("startedAt"), usage.get("endedAt"))
        if duration is not None:
            row["durationMs"] = duration


def _worker_detail_label(worker: dict) -> str:
    role = str(worker.get("role") or worker.get("workerId") or "Worker").strip()
    agent = str(worker.get("agent") or "").strip()
    status = str(worker.get("status") or "").strip()
    suffix = ", ".join(value for value in (agent, status) if value)
    return f"{role} ({suffix})" if suffix else role


def _worker_detail_row(label: str, usage: dict) -> dict:
    if usage.get("source") == "unavailable":
        return {
            "label": label,
            "totalTokens": None,
            "cacheReadTokens": None,
            "billableTokens": None,
            "costUsd": None,
            "cliTotalTokens": None,
            "cliCostUsd": None,
        }
    return {
        "label": label,
        "totalTokens": usage.get("totalTokens"),
        "cacheReadTokens": usage.get("cacheReadTokens"),
        "billableTokens": usage.get("billableEquivalentTokens"),
        "costUsd": usage.get("estimatedCostUsd"),
        "cliTotalTokens": usage.get("cliTotalTokens"),
        "cliCostUsd": usage.get("cliEstimatedCostUsd"),
    }


def _cache_reads(summary: dict, team_state: dict) -> dict[str, int]:
    """Lead / worker / grand cache-read totals for this run.

    The figure joined `usageSummary` after most team-states on disk were
    written, and re-rendering one of those reports is exactly when the reader
    goes looking for the corrected table — so a state without the summary keys
    is added up from the usage blocks it does carry.
    """
    lead = summary.get("leadCacheReadTokens")
    worker = summary.get("workerCacheReadTokens")
    if isinstance(lead, int) and isinstance(worker, int):
        return {"lead": lead, "worker": worker, "grand": lead + worker}
    lead = _cache_read_tokens(team_state.get("leadUsage"))
    worker = sum(
        _cache_read_tokens(entry.get("usage"))
        for entry in team_state.get("workers") or []
        if isinstance(entry, dict)
    )
    worker += _cache_read_tokens(summary.get("unattributedWorkerUsage"))
    return {"lead": lead, "worker": worker, "grand": lead + worker}


def _cache_read_tokens(usage: object) -> int:
    if not isinstance(usage, dict):
        return 0
    value = usage.get("cacheReadTokens")
    return value if isinstance(value, int) else 0


def _worker_detail_rows(team_state: dict) -> list[dict]:
    rows = []
    for worker in team_state.get("workers") or []:
        if isinstance(worker, dict):
            rows.append(
                _worker_detail_row(_worker_detail_label(worker), worker.get("usage") or {})
            )
    unattributed = (team_state.get("usageSummary") or {}).get("unattributedWorkerUsage")
    if isinstance(unattributed, dict):
        rows.append(_worker_detail_row("Unattributed worker usage", unattributed))
    return rows


def populate_data_token_cells(data_path: Path, team_state: dict) -> int:
    """Mutate ``data_path`` (final-report data.json) in place: fill
    ``tokenUsage`` rows and ``executionStatus`` row token cells from
    ``team_state['usageSummary']`` and ``team_state['workers']`` /
    ``team_state['lead']``. Returns the number of cells changed (rough
    count for diagnostics).

    Raises ``SubstituteRefusedError`` when ``usageSummary.grandTotalTokens``
    is zero — substituting zeros bakes in the "collector ran but found
    nothing" failure mode. Callers that intentionally want zeros (e.g.
    unit-test fixtures) must omit ``usageSummary`` from the team-state.
    """
    summary = team_state.get("usageSummary") or {}
    grand_total = summary.get("grandTotalTokens", 0)
    if isinstance(grand_total, (int, float)) and grand_total == 0 and summary:
        raise SubstituteRefusedError(
            "Refusing to substitute zero-only usageSummary into "
            f"{data_path}. grandTotalTokens=0 means the collector ran "
            "but every session jsonl was empty (or absent). Re-run "
            "after locating the missing session jsonls. To intentionally "
            "ship zeros (test fixtures only), omit `usageSummary` from "
            "team-state before calling this."
        )

    if not data_path.is_file():
        raise SubstituteRefusedError(f"data file not found: {data_path}")

    data = json.loads(data_path.read_text(encoding="utf-8"))
    changes = 0

    # Token Usage Summary table — four rows.
    cost = summary.get("estimatedCostUsd") or {}
    lead_cost = cost.get("lead") or 0
    worker_cost = cost.get("claudeWorkers") or 0
    cli_cost = cost.get("cliWorkers") or 0

    cache_reads = _cache_reads(summary, team_state)
    token_usage = data.setdefault("tokenUsage", {})
    token_usage.setdefault("lead", {}).update({
        "totalTokens": summary.get("leadTotalTokens"),
        "cacheReadTokens": cache_reads["lead"],
        "billableTokens": summary.get("leadBillableEquivalentTokens"),
        "costUsd": lead_cost,
    })
    token_usage.setdefault("worker", {}).update({
        "totalTokens": summary.get("workerTotalTokens"),
        "cacheReadTokens": cache_reads["worker"],
        "billableTokens": summary.get("workerBillableEquivalentTokens"),
        "costUsd": worker_cost,
    })
    token_usage.setdefault("grand", {}).update({
        "totalTokens": summary.get("grandTotalTokens"),
        "cacheReadTokens": cache_reads["grand"],
        "billableTokens": summary.get("grandBillableEquivalentTokens"),
        # CLI tracked on its own row per the template — grand here means
        # lead + claudeWorkers, not lead + claudeWorkers + cli.
        "costUsd": lead_cost + worker_cost,
    })
    token_usage.setdefault("cli", {})["costUsd"] = cli_cost
    token_usage["workerDetails"] = _worker_detail_rows(team_state)
    cumulative = task_cumulative_usage(data_path)
    if cumulative is not None:
        token_usage["taskCumulative"] = cumulative
    changes += 4

    # Execution Status by Agent — per-row token / cost / duration.
    # `collect.py` writes lead usage to `team_state["leadUsage"]` (flat
    # usage_block), while worker usage lives at `worker["usage"]`. Wrap
    # the lead block so `_populate_execution_row` can read both shapes
    # via the same `source["usage"]` path.
    lead_source = {"usage": team_state.get("leadUsage") or {}}
    workers = team_state.get("workers") or []
    rows = data.get("executionStatus") or []
    used_workers: set[int] = set()
    for row in rows:
        role = row.get("role", "")
        if "lead" in role.lower():
            _populate_execution_row(row, lead_source)
            changes += 1
            continue
        idx = _match_worker_index(row, workers, used_workers)
        if idx is not None:
            used_workers.add(idx)
            _populate_execution_row(row, workers[idx])
            changes += 1

    data_path.write_text(
        json.dumps(data, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    return changes


def populate_token_cells(
    data_path: Path,
    team_state: dict,
) -> int:
    """Mutate data.json token cells and materialize design-prep requests.

    The full reading copy is rendered on demand with
    ``okstra render-final-report``.
    """
    changes = populate_data_token_cells(data_path, team_state)
    if data_path.name.startswith("final-report-implementation-planning-"):
        try:
            materialize_design_prep_requests(data_path)
        except DesignPrepError as exc:
            raise SubstituteRefusedError(
                f"design-prep request materialization failed: {exc}"
            ) from exc
    return changes