"""Suite-level cost envelope evaluation for prompt eval suites.

Aggregates simulated token usage across all prompt eval cases and validates
that the total stays within a configured cost envelope.
"""
from __future__ import annotations

import json
from dataclasses import asdict, dataclass, field
from typing import Any


@dataclass(frozen=True)
class CostEnvelopeSpec:
    """Cost envelope configuration for a suite run."""

    max_total_tokens: int
    max_tokens_per_case: int | None = None
    warn_threshold_pct: float = 0.80


@dataclass(frozen=True)
class CaseCostEntry:
    """Token cost record for a single eval case."""

    case_id: str
    tokens_used: int
    over_case_limit: bool = False


@dataclass
class SuiteCostReport:
    """Aggregated cost report for an eval suite run."""

    suite_id: str
    total_tokens: int = 0
    max_total_tokens: int = 0
    case_costs: list[CaseCostEntry] = field(default_factory=list)
    status: str = "PASS"
    violations: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        """Serialize to a dictionary with schema version."""
        utilization = (
            self.total_tokens / self.max_total_tokens
            if self.max_total_tokens > 0
            else 0.0
        )
        return {
            "schema_version": "suite-cost-report/1.0",
            "suite_id": self.suite_id,
            "total_tokens": self.total_tokens,
            "max_total_tokens": self.max_total_tokens,
            "utilization_pct": round(utilization, 6),
            "case_costs": [asdict(c) for c in self.case_costs],
            "status": self.status,
            "violations": list(self.violations),
        }

    def to_json(self) -> str:
        """Serialize to deterministic JSON string."""
        return json.dumps(self.to_dict(), sort_keys=True, ensure_ascii=False)


def validate_suite_cost_envelope(
    suite_results: dict[str, Any],
    envelope: CostEnvelopeSpec,
    case_tokens: dict[str, int] | None = None,
) -> SuiteCostReport:
    """Validate that suite token usage stays within the cost envelope.

    Parameters
    ----------
    suite_results:
        An eval-result/1.0 dict with ``results`` list and ``suite_id``.
    envelope:
        The cost envelope specification (max tokens, thresholds).
    case_tokens:
        Optional explicit token counts keyed by case_id.  When a case_id
        is not present here, tokens are estimated from the serialized
        ``actual`` field (``len(json.dumps(actual)) // 4``).

    Returns
    -------
    SuiteCostReport
    """
    case_tokens = case_tokens or {}
    report = SuiteCostReport(
        suite_id=suite_results.get("suite_id", "unknown"),
        max_total_tokens=envelope.max_total_tokens,
    )

    results_sorted = sorted(
        suite_results.get("results", []),
        key=lambda r: r.get("case_id", ""),
    )

    total = 0
    for r in results_sorted:
        cid = r.get("case_id", "")
        if cid in case_tokens:
            tokens = case_tokens[cid]
        else:
            tokens = len(json.dumps(r.get("actual", ""))) // 4

        over_limit = False
        if envelope.max_tokens_per_case is not None and tokens > envelope.max_tokens_per_case:
            over_limit = True
            report.violations.append(
                f"Case {cid}: {tokens} tokens exceeds per-case limit "
                f"of {envelope.max_tokens_per_case}"
            )
            report.status = "FAIL"

        report.case_costs.append(CaseCostEntry(
            case_id=cid, tokens_used=tokens, over_case_limit=over_limit,
        ))
        total += tokens

    report.total_tokens = total

    if total > envelope.max_total_tokens:
        report.violations.append(
            f"Total {total} tokens exceeds budget of {envelope.max_total_tokens}"
        )
        report.status = "FAIL"
    elif report.status != "FAIL" and envelope.max_total_tokens > 0:
        utilization = total / envelope.max_total_tokens
        if utilization >= envelope.warn_threshold_pct:
            report.status = "WARN"

    return report
