"""Turn trial records into verdicts (docs/08-surfaces.md).

This is the layer between the raw JSONL log and anything that renders. It is
kept separate from rendering so that `ablate report` can re-analyze a stored
run with different margins at zero cost, and so the numbers can be tested
without parsing terminal output.

Metric choice follows ADR-0010: **tokens and turns are primary, pass rate is
the guardrail.** At an affordable N the token metric is comfortably detectable
and the pass-rate metric is not — by roughly an order of magnitude in the ratio
of effect size to MDE. The report shows both with equal visual weight anyway,
because a user deciding whether to delete something needs both side by side.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field

from ablate.record import TrialRecord
from ablate.stats import (
    DEFAULT_MARGIN_CONTINUOUS,
    DEFAULT_MARGIN_PASS_RATE,
    MetricVerdict,
    PairedSample,
    bootstrap_delta,
    classify,
    mde_pass_rate,
)


def _total_tokens(r: TrialRecord) -> int:
    t = r.telemetry
    return t.input_tokens + t.output_tokens + t.cache_read + t.cache_write


def _tokens_per_turn(r: TrialRecord) -> float | None:
    """None when the agent did not report token usage.

    A trial that actually ran always consumed tokens, so an all-zero count
    means the adapter had nothing to report — not that the trial was free.
    Returning 0.0 would put a fabricated number in the primary metric
    (ADR-0010) and make an unmeasurable arm look like a spectacularly cheap
    one. Codex may be exactly this case; see Q1.
    """
    total = _total_tokens(r)
    if total == 0 or r.telemetry.turns <= 0:
        return None
    return total / r.telemetry.turns


@dataclass(frozen=True)
class MetricSpec:
    """How to pull one metric out of a trial and how to judge it."""

    key: str
    label: str
    extract: Callable[[TrialRecord], float | None]
    higher_is_better: bool
    relative: bool
    margin: float
    fmt: Callable[[float], str]


#: Wall-clock time is deliberately absent as a verdict metric. It is
#: contaminated by provider load, rate-limit backoff and local machine state —
#: none of which are properties of the config (ADR-0010).
METRICS: tuple[MetricSpec, ...] = (
    MetricSpec(
        key="pass_rate",
        label="PASS RATE",
        extract=lambda r: float(r.passed),
        higher_is_better=True,
        relative=False,
        margin=DEFAULT_MARGIN_PASS_RATE,
        fmt=lambda v: f"{v:.2f}",
    ),
    MetricSpec(
        key="tokens_per_turn",
        label="TOKENS PER TURN",
        extract=_tokens_per_turn,
        higher_is_better=False,
        relative=True,
        margin=DEFAULT_MARGIN_CONTINUOUS,
        fmt=lambda v: f"{v:,.0f}",
    ),
    MetricSpec(
        key="turns",
        label="TURNS TO COMPLETION",
        extract=lambda r: float(r.telemetry.turns) if r.telemetry.turns > 0 else None,
        higher_is_better=False,
        relative=True,
        margin=DEFAULT_MARGIN_CONTINUOUS,
        fmt=lambda v: f"{v:.1f}",
    ),
)


@dataclass(frozen=True)
class MetricAnalysis:
    spec: MetricSpec
    #: `None` for an arm that reported nothing — never 0.0. See `_arm_value`.
    arm_values: dict[str, float | None]
    verdict: MetricVerdict

    @property
    def label(self) -> str:
        return self.spec.label

    @property
    def measured(self) -> bool:
        """Whether any task reported this metric in both arms.

        The one question every renderer has to ask before printing a number.
        Checking it here, rather than each caller testing for NaN, is what
        stopped "+nanpp" and "0.00" reaching four different output formats.
        """
        return self.verdict.result.n_tasks > 0


@dataclass(frozen=True)
class TaskRow:
    """One descriptive per-task row. No p-value, ever.

    Per-task deltas are descriptive by design (docs/05-statistics.md): selecting
    the largest of 20 noisy per-task deltas and calling it a finding is the
    classic failure, and the report makes it structurally impossible by never
    emitting per-task inference.
    """

    task_id: str
    passed: dict[str, int]
    total: dict[str, int]
    delta: float | None
    degenerate: str | None = None

    @property
    def is_degenerate(self) -> bool:
        return self.degenerate is not None


@dataclass(frozen=True)
class Analysis:
    run_id: str
    variant_ids: tuple[str, ...]
    n_tasks: int
    n_trials: int
    n_records: int
    metrics: list[MetricAnalysis]
    tasks: list[TaskRow]
    total_cost_usd: float
    cost_known: bool
    #: Mean per-trial cost with the first trial of each (task, arm) excluded.
    #: `None` when there is nothing left after the exclusion.
    steady_state_cost_usd: float | None
    alpha: float
    seed: int | None
    agent: str
    agent_version: str | None
    models: tuple[str, ...]
    errors: int
    warnings: list[str] = field(default_factory=list)

    def metric(self, key: str) -> MetricAnalysis | None:
        return next((m for m in self.metrics if m.spec.key == key), None)

    @property
    def degenerate_tasks(self) -> list[TaskRow]:
        return [t for t in self.tasks if t.is_degenerate]

    @property
    def informative_tasks(self) -> int:
        return self.n_tasks - len(self.degenerate_tasks)

    @property
    def mixes_models(self) -> bool:
        """A model version bump moves the baseline; mixing them invalidates Δ."""
        return len(self.models) > 1


def _paired_sample(
    records: list[TrialRecord],
    spec: MetricSpec,
    variant_a: str,
    variant_b: str,
) -> PairedSample:
    pairs: dict[str, tuple[list[float], list[float]]] = {}
    for r in records:
        if r.variant_id not in (variant_a, variant_b):
            continue
        value = spec.extract(r)
        if value is None:
            continue
        a, b = pairs.setdefault(r.task_id, ([], []))
        (a if r.variant_id == variant_a else b).append(value)
    return PairedSample.from_pairs(pairs)


def _arm_value(sample: PairedSample, arm: str) -> float | None:
    """Mean of per-task means — the same weighting the estimand uses.

    Not the pooled mean over trials: if the arms are ragged, a pooled mean
    would not satisfy `displayed_A − displayed_B == Δ̂`, and a report whose
    headline numbers do not reconcile is worse than no report.

    `None` when no task reported the metric, for the same reason
    `bootstrap_delta` returns NaN there: zero is a measurement. Rendering an
    arm that reported nothing as "0.00 pass rate" or "0 tokens" states
    something false, and states it in the direction that flatters whichever
    arm is being defended.
    """
    arrays = sample.a if arm == "a" else sample.b
    if not arrays:
        return None
    return float(sum(x.mean() for x in arrays) / len(arrays))


def analyze(
    records: list[TrialRecord],
    variant_ids: tuple[str, ...],
    *,
    run_id: str = "",
    margin_overrides: dict[str, float] | None = None,
    alpha: float = 0.05,
    seed: int | None = None,
    agent: str = "",
    agent_version: str | None = None,
) -> Analysis:
    """Compute every verdict and per-task row for a set of trials.

    Exactly two arms are supported: the estimand is a paired difference, and
    there is no honest way to extend that to three arms without a multiple
    comparison correction the design does not currently carry.
    """
    if len(variant_ids) != 2:
        raise ValueError(f"analysis needs exactly two arms, got {len(variant_ids)}")

    scored = [r for r in records if r.error is None]
    variant_a, variant_b = variant_ids
    overrides = margin_overrides or {}

    metrics: list[MetricAnalysis] = []
    pass_sample = _paired_sample(scored, METRICS[0], variant_a, variant_b)
    rate_a, rate_b = _arm_value(pass_sample, "a"), _arm_value(pass_sample, "b")
    n_trials = _median_replicates(scored)

    for spec in METRICS:
        sample = _paired_sample(scored, spec, variant_a, variant_b)
        result = bootstrap_delta(
            sample, relative=spec.relative, alpha=alpha, seed=seed
        )
        metrics.append(
            MetricAnalysis(
                spec=spec,
                arm_values={
                    variant_a: _arm_value(sample, "a"),
                    variant_b: _arm_value(sample, "b"),
                },
                verdict=classify(
                    result,
                    metric=spec.label.lower(),
                    margin=overrides.get(spec.key, spec.margin),
                    higher_is_better=spec.higher_is_better,
                    # Only feeds the power calculation in the INCONCLUSIVE
                    # message. With no pass-rate data there is no observed rate
                    # to plan against, so fall back to the least favourable
                    # assumption (p=0.5) rather than inventing an optimistic one.
                    p_a=rate_a if rate_a is not None else 0.5,
                    p_b=rate_b if rate_b is not None else 0.5,
                    n_trials=n_trials,
                ),
            )
        )

    tasks = _task_rows(scored, variant_a, variant_b)
    costs = [r.telemetry.cost_usd for r in records]

    analysis = Analysis(
        run_id=run_id,
        variant_ids=variant_ids,
        n_tasks=len({r.task_id for r in scored}),
        n_trials=n_trials,
        n_records=len(records),
        metrics=metrics,
        tasks=tasks,
        total_cost_usd=sum(c for c in costs if c is not None),
        cost_known=bool(costs) and all(c is not None for c in costs),
        steady_state_cost_usd=_steady_state_cost(records),
        alpha=alpha,
        seed=seed,
        agent=agent,
        agent_version=agent_version,
        models=tuple(sorted({r.model for r in scored if r.model})),
        errors=len(records) - len(scored),
    )
    analysis.warnings.extend(_warnings(analysis))
    return analysis


def _steady_state_cost(records: list[TrialRecord]) -> float | None:
    """Mean per-trial cost once the caches are warm (Q7).

    The first trial of each (task, arm) pays a prompt-cache *write*; every one
    after it hits a warm cache at roughly a tenth the input rate. Averaging all
    of them together overstates what the next run will cost, which is exactly
    the number a user is trying to budget with.

    "First" is chronological, not replicate 1 — the work queue is shuffled, so
    the trial that warms the cache is whichever happened to run first.
    """
    priced = [r for r in records if r.telemetry.cost_usd is not None]
    if not priced:
        return None

    first: dict[tuple[str, str], TrialRecord] = {}
    for record in sorted(priced, key=lambda r: r.started_at):
        first.setdefault((record.task_id, record.variant_id), record)

    warm = [r for r in priced if first.get((r.task_id, r.variant_id)) is not r]
    if not warm:
        return None
    return sum(r.telemetry.cost_usd for r in warm) / len(warm)


def _median_replicates(records: list[TrialRecord]) -> int:
    """Trials per (task, arm), from the data rather than the flag.

    A run stopped by the budget ceiling has fewer than requested, and every
    downstream power calculation should reflect what was actually collected.
    """
    counts: dict[tuple[str, str], int] = {}
    for r in records:
        counts[(r.task_id, r.variant_id)] = counts.get((r.task_id, r.variant_id), 0) + 1
    if not counts:
        return 0
    ordered = sorted(counts.values())
    return ordered[len(ordered) // 2]


def _task_rows(records: list[TrialRecord], variant_a: str, variant_b: str) -> list[TaskRow]:
    passed: dict[str, dict[str, int]] = {}
    total: dict[str, dict[str, int]] = {}

    for r in records:
        passed.setdefault(r.task_id, {}).setdefault(r.variant_id, 0)
        total.setdefault(r.task_id, {}).setdefault(r.variant_id, 0)
        total[r.task_id][r.variant_id] += 1
        passed[r.task_id][r.variant_id] += int(r.passed)

    rows = []
    for task_id in sorted(total):
        p, t = passed[task_id], total[task_id]
        na, nb = t.get(variant_a, 0), t.get(variant_b, 0)
        ka, kb = p.get(variant_a, 0), p.get(variant_b, 0)

        delta = None
        degenerate = None
        if na and nb:
            rate_a, rate_b = ka / na, kb / nb
            delta = rate_a - rate_b
            # A task that every trial passes in both arms, or that no trial
            # passes in either, contributes d_i = 0 with near-zero variance:
            # full price, no information, and it attenuates Δ̂ toward zero.
            if ka == na and kb == nb:
                degenerate = "all pass both arms"
            elif ka == 0 and kb == 0:
                degenerate = "none pass either arm"

        rows.append(
            TaskRow(
                task_id=task_id,
                passed={variant_a: ka, variant_b: kb},
                total={variant_a: na, variant_b: nb},
                delta=delta,
                degenerate=degenerate,
            )
        )
    return rows


def _warnings(analysis: Analysis) -> list[str]:
    out: list[str] = []

    degenerate = analysis.degenerate_tasks
    if degenerate:
        informative = analysis.informative_tasks
        message = (
            f"{len(degenerate)} of {analysis.n_tasks} tasks contributed no signal. "
            f"This suite is effectively {informative} task"
            f"{'s' if informative != 1 else ''} for this comparison"
        )
        # Quantify the cost of the degeneracy honestly: what the MDE is now
        # versus what it would be if those slots held informative tasks. No
        # claim is made about *dropping* them — excluding tasks on the basis of
        # their observed outcome is a forking path that biases Δ̂.
        if informative > 0 and analysis.n_trials > 0:
            now = mde_pass_rate(informative, analysis.n_trials)
            replaced = mde_pass_rate(analysis.n_tasks, analysis.n_trials)
            message += (
                f" — MDE {now * 100:.0f}pp, versus {replaced * 100:.0f}pp if they were "
                f"replaced with mid-difficulty tasks at the same cost"
            )
        out.append(message + ".")

    if analysis.mixes_models:
        out.append(
            f"Trials span more than one model ({', '.join(analysis.models)}). "
            "A model version bump moves the baseline, so this comparison mixes "
            "two different experiments and Δ is not interpretable."
        )

    if analysis.errors:
        out.append(
            f"{analysis.errors} trial{'s' if analysis.errors != 1 else ''} errored and "
            "were excluded from the analysis."
        )

    return out
