# platform/deskclaw/experiment/scoring.py
# ═══════════════════════════════════════════════════════════════
# MetadataScorer — Operational metrics extraction & aggregation
# ═══════════════════════════════════════════════════════════════

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from deskclaw.experiment.models import ExperimentResult, RunResult


@dataclass(frozen=True)
class RunMetrics:
    blueprint_name: str
    input_tokens: int = 0
    output_tokens: int = 0
    total_tokens: int = 0
    api_calls: int = 0
    latency_ms: int = 0
    cost_usd: float = 0.0

    def to_dict(self) -> dict[str, Any]:
        return {
            "blueprint_name": self.blueprint_name,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "total_tokens": self.total_tokens,
            "api_calls": self.api_calls,
            "latency_ms": self.latency_ms,
            "cost_usd": round(self.cost_usd, 6),
        }


@dataclass(frozen=True)
class BlueprintMetrics:
    blueprint_name: str
    run_count: int = 0
    mean_latency_ms: float = 0.0
    median_latency_ms: float = 0.0
    p95_latency_ms: float = 0.0
    p99_latency_ms: float = 0.0
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    total_tokens: int = 0
    mean_tokens_per_run: float = 0.0
    total_cost_usd: float = 0.0
    mean_cost_per_run: float = 0.0
    pass_rate: float = 0.0

    def to_dict(self) -> dict[str, Any]:
        return {
            "blueprint_name": self.blueprint_name,
            "run_count": self.run_count,
            "mean_latency_ms": round(self.mean_latency_ms, 2),
            "median_latency_ms": round(self.median_latency_ms, 2),
            "p95_latency_ms": round(self.p95_latency_ms, 2),
            "p99_latency_ms": round(self.p99_latency_ms, 2),
            "total_input_tokens": self.total_input_tokens,
            "total_output_tokens": self.total_output_tokens,
            "total_tokens": self.total_tokens,
            "mean_tokens_per_run": round(self.mean_tokens_per_run, 2),
            "total_cost_usd": round(self.total_cost_usd, 6),
            "mean_cost_per_run": round(self.mean_cost_per_run, 6),
            "pass_rate": round(self.pass_rate, 4),
        }


@dataclass
class ExperimentMetrics:
    blueprint_metrics: dict[str, BlueprintMetrics] = field(default_factory=dict)
    total_runs: int = 0
    total_tokens: int = 0
    total_cost_usd: float = 0.0
    overall_pass_rate: float = 0.0

    def to_dict(self) -> dict[str, Any]:
        return {
            "blueprint_metrics": {
                k: v.to_dict() for k, v in self.blueprint_metrics.items()
            },
            "total_runs": self.total_runs,
            "total_tokens": self.total_tokens,
            "total_cost_usd": round(self.total_cost_usd, 6),
            "overall_pass_rate": round(self.overall_pass_rate, 4),
        }


class MetadataScorer:
    def score_run(self, run: RunResult) -> RunMetrics:
        total = run.input_tokens + run.output_tokens
        return RunMetrics(
            blueprint_name=run.blueprint_name,
            input_tokens=run.input_tokens,
            output_tokens=run.output_tokens,
            total_tokens=total,
            api_calls=run.api_calls,
            latency_ms=run.elapsed_ms,
            cost_usd=run.cost_usd,
        )

    def score_experiment(self, result: ExperimentResult) -> ExperimentMetrics:
        blueprint_names = {r.blueprint_name for r in result.runs}
        bp_metrics: dict[str, BlueprintMetrics] = {}

        for bp_name in sorted(blueprint_names):
            runs = result.runs_for_blueprint(bp_name)
            bp_metrics[bp_name] = self._aggregate_blueprint(bp_name, runs)

        total_runs = len(result.runs)
        total_tokens = sum(m.total_tokens for m in bp_metrics.values())
        total_cost = sum(m.total_cost_usd for m in bp_metrics.values())
        passed = sum(1 for r in result.runs if r.status == "passed")
        pass_rate = passed / max(total_runs, 1)

        return ExperimentMetrics(
            blueprint_metrics=bp_metrics,
            total_runs=total_runs,
            total_tokens=total_tokens,
            total_cost_usd=total_cost,
            overall_pass_rate=pass_rate,
        )

    def _aggregate_blueprint(
        self, bp_name: str, runs: list[RunResult]
    ) -> BlueprintMetrics:
        if not runs:
            return BlueprintMetrics(blueprint_name=bp_name)

        latencies = [r.elapsed_ms for r in runs]
        input_tok = sum(r.input_tokens for r in runs)
        output_tok = sum(r.output_tokens for r in runs)
        total_tok = input_tok + output_tok
        costs = sum(r.cost_usd for r in runs)
        passed = sum(1 for r in runs if r.status == "passed")
        n = len(runs)

        return BlueprintMetrics(
            blueprint_name=bp_name,
            run_count=n,
            mean_latency_ms=sum(latencies) / n,
            median_latency_ms=_percentile(latencies, 0.50),
            p95_latency_ms=_percentile(latencies, 0.95),
            p99_latency_ms=_percentile(latencies, 0.99),
            total_input_tokens=input_tok,
            total_output_tokens=output_tok,
            total_tokens=total_tok,
            mean_tokens_per_run=total_tok / n,
            total_cost_usd=costs,
            mean_cost_per_run=costs / n,
            pass_rate=passed / n,
        )


def _percentile(values: list[int] | list[float], p: float) -> float:
    if not values:
        return 0.0
    sorted_vals = sorted(values)
    idx = int(p * (len(sorted_vals) - 1))
    return float(sorted_vals[idx])
