"""Tradeoffs analysis: cost, quality, and latency optimization.

Deterministic — no LLM required. Identifies where you can get
the same quality for less money or time.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from bandito.models import get_eval_score
from ..review.aggregator import aggregate_summaries
from ..review.models import AggregateStats, ModelBreakdown, Segment
from ..review.segmenter import segment
from ..review.utils import build_trace_model_map, get_primary_model


@dataclass
class ModelEfficiency:
    """Per-model quality/cost/latency comparison."""
    model: str
    provider: str | None
    trace_count: int
    quality: float | None  # judge or human avg grade, None if no evals
    avg_cost: float | None  # None if no cost data
    avg_latency: float | None
    total_cost: float | None
    has_cost: bool = False
    annotation: str = ""  # "← best quality/$", "← cheapest", etc.


@dataclass
class SegmentOpportunity:
    """A segment where model switching could save cost."""
    segment_key: str
    trace_count: int
    models: list[ModelEfficiency]
    insight: str  # deterministic one-liner


@dataclass
class CostDriver:
    """What's driving cost."""
    label: str
    detail: str


@dataclass
class QualityVsCost:
    """Cheap half vs expensive half comparison."""
    cheap_quality: float
    cheap_cost: float
    expensive_quality: float
    expensive_cost: float
    cost_ratio: float  # how many x more expensive
    quality_gap: float  # percentage point difference
    has_data: bool = False


@dataclass
class TradeoffsResult:
    """Full tradeoffs analysis output."""
    project: str
    trace_count: int
    model_count: int
    segments_with_evals: int
    has_evaluations: bool

    models: list[ModelEfficiency] = field(default_factory=list)
    segment_opportunities: list[SegmentOpportunity] = field(default_factory=list)
    cost_drivers: list[CostDriver] = field(default_factory=list)
    quality_vs_cost: QualityVsCost = field(default_factory=lambda: QualityVsCost(0, 0, 0, 0, 0, 0))

    def to_dict(self) -> dict[str, Any]:
        return {
            "project": self.project,
            "trace_count": self.trace_count,
            "model_count": self.model_count,
            "has_evaluations": self.has_evaluations,
            "models": [
                {
                    "model": m.model,
                    "provider": m.provider,
                    "trace_count": m.trace_count,
                    "quality": m.quality,
                    "avg_cost": m.avg_cost,
                    "avg_latency": round(m.avg_latency, 3) if m.avg_latency else None,
                    "total_cost": m.total_cost,
                    "annotation": m.annotation,
                }
                for m in self.models
            ],
            "segment_opportunities": [
                {
                    "segment": s.segment_key,
                    "trace_count": s.trace_count,
                    "insight": s.insight,
                    "models": [
                        {
                            "model": m.model,
                            "quality": m.quality,
                            "avg_cost": m.avg_cost,
                        }
                        for m in s.models
                    ],
                }
                for s in self.segment_opportunities
            ],
            "cost_drivers": [
                {"label": d.label, "detail": d.detail}
                for d in self.cost_drivers
            ],
            "quality_vs_cost": {
                "has_data": self.quality_vs_cost.has_data,
                "cheap_half": {
                    "quality": self.quality_vs_cost.cheap_quality,
                    "avg_cost": self.quality_vs_cost.cheap_cost,
                },
                "expensive_half": {
                    "quality": self.quality_vs_cost.expensive_quality,
                    "avg_cost": self.quality_vs_cost.expensive_cost,
                },
                "cost_ratio": round(self.quality_vs_cost.cost_ratio, 1),
                "quality_gap_pp": round(self.quality_vs_cost.quality_gap * 100, 1),
            } if self.quality_vs_cost.has_data else None,
        }


def run_tradeoffs(
    project: str,
    traces: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]] | None = None,
) -> TradeoffsResult:
    """Run tradeoffs analysis on traces and evaluations."""
    evaluations = evaluations or {}

    stats = aggregate_summaries(traces, evaluations)
    segments = segment(traces, evaluations, from_summaries=True)

    # Model efficiency
    model_efficiencies = _compute_model_efficiency(stats, traces, evaluations)

    # Segment opportunities (only where multiple models exist within a segment)
    seg_opportunities = _compute_segment_opportunities(segments, traces, evaluations)

    # Cost drivers
    cost_drivers = _compute_cost_drivers(stats, traces)

    # Quality vs cost
    qvc = _compute_quality_vs_cost(traces, evaluations)

    segments_with_evals = sum(
        1 for s in segments
        if s.stats.judge_eval_count > 0 or s.stats.human_eval_count > 0
    )

    return TradeoffsResult(
        project=project,
        trace_count=stats.trace_count,
        model_count=len(stats.models),
        segments_with_evals=segments_with_evals,
        has_evaluations=stats.eval_count > 0,
        models=model_efficiencies,
        segment_opportunities=seg_opportunities,
        cost_drivers=cost_drivers,
        quality_vs_cost=qvc,
    )


def _compute_model_efficiency(
    stats: AggregateStats,
    traces: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]],
) -> list[ModelEfficiency]:
    """Build per-model quality/cost/latency comparison."""
    if not stats.models:
        return []

    trace_model = build_trace_model_map(traces)

    # Aggregate quality per model
    model_scores: dict[str, list[float]] = {}
    for trace_id, ev in evaluations.items():
        model = trace_model.get(trace_id)
        if not model:
            continue
        score = get_eval_score(ev) if ev else None
        if score is not None:
            model_scores.setdefault(model, []).append(score)

    # Count traces per model
    model_trace_counts: dict[str, int] = {}
    for model in trace_model.values():
        model_trace_counts[model] = model_trace_counts.get(model, 0) + 1

    # Compute per-model cost from trace-level data (works for summaries too)
    model_costs: dict[str, list[float]] = {}
    for t in traces:
        model = trace_model.get(t.get("trace_id", ""))
        if not model:
            continue
        cost = t.get("total_cost")
        if cost is None and "spans" in t:
            cost = sum(sp.get("cost") or 0 for sp in t["spans"] if sp.get("cost") is not None)
            if cost == 0:
                cost = None
        if cost is not None:
            model_costs.setdefault(model, []).append(cost)

    efficiencies = []
    for mb in stats.models:
        scores = model_scores.get(mb.model, [])
        quality = sum(scores) / len(scores) if scores else None
        costs = model_costs.get(mb.model, [])
        has_cost = len(costs) > 0 or mb.has_cost
        avg_cost = sum(costs) / len(costs) if costs else (mb.avg_cost if mb.has_cost else None)
        total_cost = sum(costs) if costs else (mb.total_cost if mb.has_cost else None)
        efficiencies.append(ModelEfficiency(
            model=mb.model,
            provider=mb.provider,
            trace_count=model_trace_counts.get(mb.model, mb.count),
            quality=quality,
            avg_cost=avg_cost if has_cost else None,
            avg_latency=mb.avg_latency if mb.avg_latency > 0 else None,
            total_cost=total_cost if has_cost else None,
            has_cost=has_cost,
        ))

    # Annotate
    _annotate_models(efficiencies)
    return efficiencies


def _annotate_models(models: list[ModelEfficiency]) -> None:
    """Add one-liner annotations to model efficiencies."""
    if len(models) < 2:
        return

    # Best quality/$ (highest quality per dollar)
    with_cost_and_quality = [
        m for m in models if m.has_cost and m.avg_cost and m.avg_cost > 0 and m.quality is not None
    ]
    if with_cost_and_quality:
        best_qpd = max(with_cost_and_quality, key=lambda m: m.quality / m.avg_cost)
        cheapest = min(with_cost_and_quality, key=lambda m: m.avg_cost)

        if best_qpd is cheapest:
            best_qpd.annotation = "\u2190 best quality/$ and cheapest"
        else:
            best_qpd.annotation = "\u2190 best quality/$"
            cheapest.annotation = "\u2190 cheapest"
    elif any(m.has_cost for m in models):
        cheapest = min([m for m in models if m.has_cost], key=lambda m: m.avg_cost)
        cheapest.annotation = "\u2190 cheapest"


def _compute_segment_opportunities(
    segments: list[Segment],
    traces: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]],
) -> list[SegmentOpportunity]:
    """Find segments where model switching could save cost."""
    opportunities = []

    trace_by_id = {t["trace_id"]: t for t in traces if t.get("trace_id")}
    trace_model = build_trace_model_map(traces)

    for seg in segments:
        if len(seg.trace_ids) < 3:
            continue

        # Find models used within this segment
        model_traces: dict[str, list[str]] = {}
        for tid in seg.trace_ids:
            model = trace_model.get(tid)
            if model:
                model_traces.setdefault(model, []).append(tid)

        if len(model_traces) < 2:
            continue

        # Compute per-model quality within this segment
        seg_models = []
        for model, tids in sorted(model_traces.items(), key=lambda x: len(x[1]), reverse=True):
            scores = []
            costs = []
            for tid in tids:
                ev = evaluations.get(tid)
                s = get_eval_score(ev) if ev else None
                if s is not None:
                    scores.append(s)
                t = trace_by_id.get(tid)
                if t:
                    c = t.get("total_cost")
                    if c is None and "spans" in t:
                        c = sum(sp.get("cost") or 0 for sp in t["spans"])
                    if c and c > 0:
                        costs.append(c)

            quality = sum(scores) / len(scores) if scores else None
            avg_cost = sum(costs) / len(costs) if costs else None

            seg_models.append(ModelEfficiency(
                model=model,
                provider=None,
                trace_count=len(tids),
                quality=quality,
                avg_cost=avg_cost,
                avg_latency=None,
                total_cost=None,
                has_cost=bool(costs),
            ))

        # Generate insight if there's a meaningful comparison
        insight = _generate_segment_insight(seg_models)
        if insight:
            opportunities.append(SegmentOpportunity(
                segment_key=seg.key,
                trace_count=seg.trace_count,
                models=seg_models,
                insight=insight,
            ))

    return opportunities


def _generate_segment_insight(models: list[ModelEfficiency]) -> str:
    """Generate a one-liner insight comparing models within a segment."""
    with_quality = [m for m in models if m.quality is not None and m.has_cost and m.avg_cost]
    if len(with_quality) < 2:
        return ""

    # Find cheapest and most expensive with quality
    cheapest = min(with_quality, key=lambda m: m.avg_cost)
    most_expensive = max(with_quality, key=lambda m: m.avg_cost)

    if cheapest is most_expensive:
        return ""

    quality_gap = most_expensive.quality - cheapest.quality
    cost_ratio = most_expensive.avg_cost / cheapest.avg_cost if cheapest.avg_cost > 0 else 0

    if abs(quality_gap) <= 0.05 and cost_ratio > 1.5:
        savings = (1 - cheapest.avg_cost / most_expensive.avg_cost) * 100
        return f"{quality_gap:+.0%} quality gap, {savings:.0f}% cost savings with {cheapest.model}"
    elif quality_gap > 0.1:
        return f"{most_expensive.model} scores {quality_gap:.0%} higher — worth the {cost_ratio:.1f}x cost"
    elif quality_gap < -0.05:
        return f"{cheapest.model} is cheaper AND scores higher"

    return ""


def _compute_cost_drivers(
    stats: AggregateStats,
    traces: list[dict[str, Any]],
) -> list[CostDriver]:
    """Identify what's driving cost."""
    drivers = []

    if not stats.cost.count:
        return drivers

    # By model
    models_with_cost = [m for m in stats.models if m.has_cost and m.total_cost > 0]
    if models_with_cost and stats.cost.total > 0:
        top_model = max(models_with_cost, key=lambda m: m.total_cost)
        pct = top_model.total_cost / stats.cost.total * 100
        if pct > 50:
            drivers.append(CostDriver(
                label="Model concentration",
                detail=(
                    f"{top_model.model} accounts for {pct:.0f}% of total cost "
                    f"(${top_model.total_cost:.4f} of ${stats.cost.total:.4f})"
                ),
            ))

    # By span kind (only available from full traces)
    has_spans = traces and "spans" in traces[0]
    if has_spans:
        kind_costs: dict[str, float] = {}
        for t in traces:
            for span in t.get("spans", []):
                cost = span.get("cost")
                if cost is not None:
                    kind = span.get("kind", "unknown")
                    kind_costs[kind] = kind_costs.get(kind, 0) + cost
        total_kind_cost = sum(kind_costs.values())
        if total_kind_cost > 0:
            for kind, cost in sorted(kind_costs.items(), key=lambda x: x[1], reverse=True):
                pct = cost / total_kind_cost * 100
                if pct > 10:
                    drivers.append(CostDriver(
                        label=f"{kind} spans",
                        detail=f"{pct:.0f}% of cost (${cost:.4f})",
                    ))

    # Outliers
    if stats.trace_count >= 10:
        trace_costs = []
        for t in traces:
            c = t.get("total_cost")
            if c is None and "spans" in t:
                c = sum(sp.get("cost") or 0 for sp in t["spans"])
            if c and c > 0:
                trace_costs.append((t.get("trace_id", "?"), c))
        if trace_costs:
            trace_costs.sort(key=lambda x: x[1], reverse=True)
            top_n = max(1, len(trace_costs) // 10)
            top_cost = sum(c for _, c in trace_costs[:top_n])
            if stats.cost.total > 0:
                ratio = top_cost / stats.cost.total
                if ratio > 0.4:
                    evidence = [f"{tid}: ${c:.4f}" for tid, c in trace_costs[:3]]
                    drivers.append(CostDriver(
                        label="Cost outliers",
                        detail=(
                            f"Top {top_n} traces account for {ratio:.0%} of cost"
                            + (f" \u00b7 " + " \u00b7 ".join(evidence) if evidence else "")
                        ),
                    ))

    return drivers


def _compute_quality_vs_cost(
    traces: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]],
) -> QualityVsCost:
    """Compare quality between cheap and expensive traces."""
    graded: list[tuple[float, float]] = []  # (cost, score)
    for t in traces:
        tid = t.get("trace_id", "")
        ev = evaluations.get(tid)
        score = get_eval_score(ev) if ev else None
        if score is None:
            continue
        # Summary path: total_cost. Full trace path: sum from spans.
        cost = t.get("total_cost")
        if cost is None and "spans" in t:
            cost = sum(sp.get("cost") or 0 for sp in t["spans"])
        if cost and cost > 0:
            graded.append((cost, score))

    if len(graded) < 10:
        return QualityVsCost(0, 0, 0, 0, 0, 0)

    graded.sort(key=lambda x: x[0])
    mid = len(graded) // 2
    cheap_scores = [s for _, s in graded[:mid]]
    expensive_scores = [s for _, s in graded[mid:]]
    cheap_costs = [c for c, _ in graded[:mid]]
    expensive_costs = [c for c, _ in graded[mid:]]

    cheap_q = sum(cheap_scores) / len(cheap_scores)
    expensive_q = sum(expensive_scores) / len(expensive_scores)
    cheap_c = sum(cheap_costs) / len(cheap_costs)
    expensive_c = sum(expensive_costs) / len(expensive_costs)

    return QualityVsCost(
        cheap_quality=cheap_q,
        cheap_cost=cheap_c,
        expensive_quality=expensive_q,
        expensive_cost=expensive_c,
        cost_ratio=expensive_c / cheap_c if cheap_c > 0 else 0,
        quality_gap=expensive_q - cheap_q,
        has_data=True,
    )
