"""Provenance-checked token baseline comparison for the v3 release gate.

The target intentionally does not accept estimated tokens or a hand-written
percentage.  A v2 run must record the same fixed requests and author request
configuration before it can be compared with a v3 release report.
"""

from __future__ import annotations

from hashlib import sha256
import json
from pathlib import Path
from statistics import median
from typing import Any

from app.settings import ProviderConfig, ProviderModel


BASELINE_SCHEMA_VERSION = "cad.token-baseline.v1"
BASELINE_PROTOCOL_VERSION = "2.0"
MINIMUM_REPETITIONS = 3
TOKEN_REDUCTION_TARGET = 0.30
MEASUREMENTS = {"prompt_tokens", "context_chars"}


class TokenBaselineError(ValueError):
    """A baseline cannot prove the release token target."""


def request_sha256(request: str) -> str:
    return sha256(request.encode("utf-8")).hexdigest()


def profile_sha256(profile_path: Path) -> str:
    try:
        return sha256(profile_path.read_bytes()).hexdigest()
    except OSError as error:
        raise TokenBaselineError("Runtime profile is unavailable for token-baseline provenance.") from error


def author_request_identity(provider: ProviderConfig, model: ProviderModel) -> dict[str, Any]:
    """Capture every author setting that can affect prompt-token comparison."""
    return {
        "provider": provider.id,
        "model": model.id,
        "api_style": provider.api_style,
        "reasoning_effort": provider.reasoning_effort,
        # Chat Completions is explicitly deterministic in StructuredModelGateway.
        # Responses models use their configured/provider-default sampling mode.
        "sampling": {"temperature": 0 if provider.api_style == "chat_completions" else None},
    }


def load_token_baseline(path: Path) -> dict[str, Any]:
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError) as error:
        raise TokenBaselineError("Token baseline report is missing or is not valid JSON.") from error
    if not isinstance(payload, dict):
        raise TokenBaselineError("Token baseline report must be a JSON object.")
    if payload.get("schema_version") != BASELINE_SCHEMA_VERSION:
        raise TokenBaselineError(f"Token baseline must use {BASELINE_SCHEMA_VERSION}.")
    if payload.get("protocol_version") != BASELINE_PROTOCOL_VERSION:
        raise TokenBaselineError("Token baseline must be a measured pre-v3 (2.0) report.")
    if payload.get("measurement") not in MEASUREMENTS:
        raise TokenBaselineError("Token baseline measurement must be prompt_tokens or context_chars.")
    return payload


def validate_token_baseline_provenance(
    baseline: dict[str, Any],
    *,
    scenarios: list[dict[str, Any]],
    author_identity: dict[str, Any],
    runtime_profile_hash: str,
) -> list[str]:
    """Reject an incomparable baseline before any billable v3 API call."""
    errors: list[str] = []
    if baseline.get("schema_version") != BASELINE_SCHEMA_VERSION:
        errors.append(f"schema_version must equal {BASELINE_SCHEMA_VERSION}")
    if baseline.get("protocol_version") != BASELINE_PROTOCOL_VERSION:
        errors.append("protocol_version must equal 2.0")
    if baseline.get("author") != author_identity:
        errors.append("author provider/model/API mode/reasoning/sampling does not match the v3 release")
    if baseline.get("runtime_profile_sha256") != runtime_profile_hash:
        errors.append("runtime profile hash does not match the v3 release")
    measurement = str(baseline.get("measurement") or "")
    if measurement not in MEASUREMENTS:
        errors.append("measurement must be prompt_tokens or context_chars")
    baseline_scenarios = baseline.get("scenarios")
    if not isinstance(baseline_scenarios, list):
        return [*errors, "scenarios must be an array"]
    by_scenario = {
        str(item.get("scenario") or ""): item
        for item in baseline_scenarios
        if isinstance(item, dict) and item.get("scenario")
    }
    for scenario in scenarios:
        scenario_id = str(scenario.get("id") or "")
        expected_request_hash = request_sha256(str(scenario.get("request") or ""))
        baseline_scenario = by_scenario.get(scenario_id)
        if not isinstance(baseline_scenario, dict):
            errors.append(f"missing baseline scenario {scenario_id}")
            continue
        if baseline_scenario.get("request_sha256") != expected_request_hash:
            errors.append(f"baseline request hash differs for {scenario_id}")
        repetitions = baseline_scenario.get("repetitions")
        if not isinstance(repetitions, list):
            errors.append(f"baseline {scenario_id} repetitions must be an array")
            continue
        by_repetition = {
            int(item.get("repetition")): item
            for item in repetitions
            if isinstance(item, dict) and isinstance(item.get("repetition"), int)
        }
        expected_repetitions = set(range(1, MINIMUM_REPETITIONS + 1))
        if not expected_repetitions.issubset(by_repetition):
            errors.append(f"baseline {scenario_id} is missing repetitions 1..{MINIMUM_REPETITIONS}")
            continue
        for repetition in expected_repetitions:
            run = by_repetition[repetition]
            if _non_negative_int(run.get("author_metric")) is None or _non_negative_int(run.get("plan_review_metric")) is None:
                errors.append(f"baseline {scenario_id} repetition {repetition} has invalid {measurement} metrics")
            for field in ("completed", "final_review_passed", "deterministic_claims_pass"):
                if not isinstance(run.get(field), bool):
                    errors.append(f"baseline {scenario_id} repetition {repetition} has non-boolean {field}")
    return errors


def compare_token_baseline(
    baseline: dict[str, Any],
    *,
    scenarios: list[dict[str, Any]],
    v3_results: list[dict[str, Any]],
    author_identity: dict[str, Any],
    runtime_profile_hash: str,
) -> dict[str, Any]:
    """Return a complete, fail-closed comparison for the release report."""
    errors = validate_token_baseline_provenance(
        baseline, scenarios=scenarios, author_identity=author_identity,
        runtime_profile_hash=runtime_profile_hash,
    )
    measurement = str(baseline.get("measurement") or "")

    baseline_scenarios = baseline.get("scenarios")
    if not isinstance(baseline_scenarios, list):
        errors.append("scenarios must be an array")
        baseline_scenarios = []
    by_scenario = {
        str(item.get("scenario") or ""): item
        for item in baseline_scenarios
        if isinstance(item, dict) and item.get("scenario")
    }
    current_by_scenario: dict[str, list[dict[str, Any]]] = {}
    for result in v3_results:
        if isinstance(result, dict) and isinstance(result.get("scenario"), str):
            current_by_scenario.setdefault(result["scenario"], []).append(result)

    baseline_tokens: list[int] = []
    current_tokens: list[int] = []
    baseline_completed: list[bool] = []
    current_completed: list[bool] = []
    baseline_final_review: list[bool] = []
    current_final_review: list[bool] = []
    baseline_deterministic: list[bool] = []
    current_deterministic: list[bool] = []
    per_scenario: list[dict[str, Any]] = []

    for scenario in scenarios:
        scenario_id = str(scenario.get("id") or "")
        expected_request_hash = request_sha256(str(scenario.get("request") or ""))
        baseline_scenario = by_scenario.get(scenario_id)
        if not isinstance(baseline_scenario, dict):
            errors.append(f"missing baseline scenario {scenario_id}")
            continue
        if baseline_scenario.get("request_sha256") != expected_request_hash:
            errors.append(f"baseline request hash differs for {scenario_id}")
            continue
        repetitions = baseline_scenario.get("repetitions")
        if not isinstance(repetitions, list) or len(repetitions) < MINIMUM_REPETITIONS:
            errors.append(f"baseline {scenario_id} needs at least {MINIMUM_REPETITIONS} repetitions")
            continue
        current = current_by_scenario.get(scenario_id, [])
        if len(current) < MINIMUM_REPETITIONS:
            errors.append(f"v3 {scenario_id} needs at least {MINIMUM_REPETITIONS} repetitions")
            continue
        baseline_by_repetition = {
            int(item.get("repetition")) : item
            for item in repetitions
            if isinstance(item, dict) and isinstance(item.get("repetition"), int)
        }
        current_by_repetition = {
            int(item.get("repetition")): item
            for item in current
            if isinstance(item.get("repetition"), int)
        }
        expected_repetitions = set(range(1, MINIMUM_REPETITIONS + 1))
        if not expected_repetitions.issubset(baseline_by_repetition):
            errors.append(f"baseline {scenario_id} is missing repetitions 1..{MINIMUM_REPETITIONS}")
            continue
        if not expected_repetitions.issubset(current_by_repetition):
            errors.append(f"v3 {scenario_id} is missing repetitions 1..{MINIMUM_REPETITIONS}")
            continue

        scenario_baseline_tokens: list[int] = []
        scenario_current_tokens: list[int] = []
        for repetition in sorted(expected_repetitions):
            baseline_run = baseline_by_repetition[repetition]
            author_tokens = _non_negative_int(baseline_run.get("author_metric"))
            plan_review_tokens = _non_negative_int(baseline_run.get("plan_review_metric"))
            if author_tokens is None or plan_review_tokens is None:
                errors.append(f"baseline {scenario_id} repetition {repetition} has invalid {measurement} metrics")
                continue
            current_run = current_by_repetition[repetition]
            usage = current_run.get("usage") if isinstance(current_run.get("usage"), dict) else {}
            v3_author_tokens = _v3_author_metric(usage, measurement)
            if v3_author_tokens is None:
                errors.append(f"v3 {scenario_id} repetition {repetition} has invalid author {measurement} usage")
                continue
            scenario_baseline_tokens.append(author_tokens + plan_review_tokens)
            scenario_current_tokens.append(v3_author_tokens)
            baseline_completed.append(bool(baseline_run.get("completed")))
            current_completed.append(str((current_run.get("projection") or {}).get("phase") or "") == "COMPLETED")
            baseline_final_review.append(bool(baseline_run.get("final_review_passed")))
            current_final_review.append(str((current_run.get("terminal") or {}).get("lifecycle") or "") == "completed")
            baseline_deterministic.append(bool(baseline_run.get("deterministic_claims_pass")))
            checks = current_run.get("checks") if isinstance(current_run.get("checks"), dict) else {}
            current_deterministic.append(bool(checks.get("deterministic_claims_pass")))
        if len(scenario_baseline_tokens) == MINIMUM_REPETITIONS and len(scenario_current_tokens) == MINIMUM_REPETITIONS:
            baseline_tokens.extend(scenario_baseline_tokens)
            current_tokens.extend(scenario_current_tokens)
            per_scenario.append({
                "scenario": scenario_id,
                "baseline_median_metric": median(scenario_baseline_tokens),
                "v3_median_author_metric": median(scenario_current_tokens),
            })

    baseline_median = median(baseline_tokens) if baseline_tokens else None
    current_median = median(current_tokens) if current_tokens else None
    reduction = (1 - current_median / baseline_median) if baseline_median and current_median is not None else None
    baseline_valid = not errors
    checks = {
        "baseline_valid": baseline_valid,
        "median_metric_reduction": baseline_valid and reduction is not None and reduction >= TOKEN_REDUCTION_TARGET,
        "completion_rate_not_lower": baseline_valid and _rate(current_completed) >= _rate(baseline_completed) if baseline_completed and current_completed else False,
        "final_review_rate_not_lower": baseline_valid and _rate(current_final_review) >= _rate(baseline_final_review) if baseline_final_review and current_final_review else False,
        "deterministic_claim_rate_not_lower": baseline_valid and _rate(current_deterministic) >= _rate(baseline_deterministic) if baseline_deterministic and current_deterministic else False,
        "plan_review_calls_removed": baseline_valid and all(
            not str(record.get("tool") or "").startswith("modeling_plan")
            for result in v3_results if isinstance(result, dict)
            for record in ((result.get("usage") or {}).get("records") or [])
            if isinstance(record, dict)
        ),
    }
    return {
        "schema_version": "cad.token-comparison.v1",
        "baseline_protocol_version": baseline.get("protocol_version"),
        "measurement": measurement,
        "baseline_median_metric": baseline_median,
        "v3_median_author_metric": current_median,
        "median_metric_reduction": reduction,
        "target_median_metric_reduction": TOKEN_REDUCTION_TARGET,
        "baseline_completion_rate": _rate(baseline_completed),
        "v3_completion_rate": _rate(current_completed),
        "baseline_final_review_rate": _rate(baseline_final_review),
        "v3_final_review_rate": _rate(current_final_review),
        "baseline_deterministic_claim_rate": _rate(baseline_deterministic),
        "v3_deterministic_claim_rate": _rate(current_deterministic),
        "per_scenario": per_scenario,
        "errors": errors,
        "checks": checks,
    }


def _non_negative_int(value: Any) -> int | None:
    return value if isinstance(value, int) and value >= 0 else None


def _v3_author_metric(usage: dict[str, Any], measurement: str) -> int | None:
    records = usage.get("records")
    if not isinstance(records, list):
        return None
    author_records = [item for item in records if isinstance(item, dict) and item.get("role") != "reviewer"]
    if not author_records:
        return None
    if measurement == "prompt_tokens" and not all(item.get("usage_available") is True for item in author_records):
        return None
    field = "prompt_tokens" if measurement == "prompt_tokens" else "context_chars"
    values = [
        _non_negative_int(item.get(field))
        for item in author_records
    ]
    return sum(value for value in values if value is not None) if all(value is not None for value in values) else None


def _rate(values: list[bool]) -> float | None:
    return sum(values) / len(values) if values else None