"""Tests for tradeoffs analysis."""

from __future__ import annotations

from datetime import datetime, timedelta, timezone
from typing import Any

import pytest

from bandito.analyze.tradeoffs import run_tradeoffs


def _make_summary(
    trace_id: str,
    model: str = "gpt-4o",
    cost: float | None = 0.01,
    latency: float = 1.0,
    tags: list[str] | None = None,
    error: bool = False,
    span_kinds: list[str] | None = None,
    total_tokens: int | None = 150,
) -> dict[str, Any]:
    """Create a trace summary for testing."""
    now = datetime.now(timezone.utc)
    start = (now - timedelta(seconds=latency)).isoformat().replace("+00:00", "Z")
    end = now.isoformat().replace("+00:00", "Z")

    return {
        "trace_id": trace_id,
        "project": "test",
        "started_at": start,
        "ended_at": end,
        "status": "error" if error else "ok",
        "tags": tags or [],
        "input_preview": "hello",
        "output_preview": "world",
        "span_count": 1,
        "span_kinds": span_kinds or ["llm"],
        "primary_model": model,
        "total_cost": cost,
        "total_tokens": total_tokens,
        "latency": latency,
        "error_count": 1 if error else 0,
    }


def _make_eval(trace_id: str, judge_score: int | None = None, human_score: int | None = None) -> dict[str, Any]:
    ev: dict[str, Any] = {"trace_id": trace_id}
    if judge_score is not None:
        ev["judge_score"] = judge_score
    if human_score is not None:
        ev["human_score"] = human_score
    return ev


class TestRunTradeoffs:
    def test_basic_result(self):
        summaries = [_make_summary(f"t-{i}") for i in range(5)]
        result = run_tradeoffs("test", summaries)
        assert result.project == "test"
        assert result.trace_count == 5
        assert result.model_count == 1
        assert not result.has_evaluations

    def test_no_traces(self):
        result = run_tradeoffs("test", [])
        assert result.trace_count == 0
        assert result.models == []

    def test_model_efficiency_without_evals(self):
        summaries = [_make_summary(f"t-{i}") for i in range(5)]
        result = run_tradeoffs("test", summaries)
        assert len(result.models) == 1
        assert result.models[0].model == "gpt-4o"
        assert result.models[0].quality is None

    def test_model_efficiency_with_evals(self):
        summaries = [_make_summary(f"t-{i}", model="gpt-4o") for i in range(5)]
        evals = {f"t-{i}": _make_eval(f"t-{i}", judge_score=1) for i in range(5)}
        result = run_tradeoffs("test", summaries, evals)
        assert result.models[0].quality == 1.0

    def test_multiple_models(self):
        summaries = (
            [_make_summary(f"a-{i}", model="gpt-4o", cost=0.02) for i in range(5)]
            + [_make_summary(f"b-{i}", model="gpt-4o-mini", cost=0.002) for i in range(5)]
        )
        evals = {
            **{f"a-{i}": _make_eval(f"a-{i}", judge_score=1) for i in range(5)},
            **{f"b-{i}": _make_eval(f"b-{i}", judge_score=1) for i in range(4)},
            "b-4": _make_eval("b-4", judge_score=0),
        }
        result = run_tradeoffs("test", summaries, evals)
        assert len(result.models) == 2
        for m in result.models:
            assert m.quality is not None

    def test_annotations_applied(self):
        summaries = (
            [_make_summary(f"a-{i}", model="gpt-4o", cost=0.02) for i in range(5)]
            + [_make_summary(f"b-{i}", model="gpt-4o-mini", cost=0.002) for i in range(5)]
        )
        evals = {
            **{f"a-{i}": _make_eval(f"a-{i}", judge_score=1) for i in range(5)},
            **{f"b-{i}": _make_eval(f"b-{i}", judge_score=1) for i in range(5)},
        }
        result = run_tradeoffs("test", summaries, evals)
        annotations = [m.annotation for m in result.models if m.annotation]
        assert len(annotations) > 0

    def test_no_cost_data(self):
        summaries = [_make_summary(f"t-{i}", cost=None) for i in range(5)]
        result = run_tradeoffs("test", summaries)
        assert len(result.models) == 1
        assert result.cost_drivers == []

    def test_cost_drivers(self):
        summaries = [_make_summary(f"t-{i}", cost=0.05) for i in range(10)]
        result = run_tradeoffs("test", summaries)
        assert result.cost_drivers is not None

    def test_quality_vs_cost(self):
        summaries = (
            [_make_summary(f"cheap-{i}", cost=0.002) for i in range(10)]
            + [_make_summary(f"expensive-{i}", cost=0.05) for i in range(10)]
        )
        evals = {
            **{f"cheap-{i}": _make_eval(f"cheap-{i}", judge_score=1) for i in range(10)},
            **{f"expensive-{i}": _make_eval(f"expensive-{i}", judge_score=1) for i in range(10)},
        }
        result = run_tradeoffs("test", summaries, evals)
        assert result.quality_vs_cost.has_data
        assert result.quality_vs_cost.cheap_cost < result.quality_vs_cost.expensive_cost

    def test_quality_vs_cost_needs_min_traces(self):
        summaries = [_make_summary(f"t-{i}", cost=0.01) for i in range(5)]
        evals = {f"t-{i}": _make_eval(f"t-{i}", judge_score=1) for i in range(5)}
        result = run_tradeoffs("test", summaries, evals)
        assert not result.quality_vs_cost.has_data

    def test_to_dict(self):
        summaries = [_make_summary(f"t-{i}") for i in range(5)]
        result = run_tradeoffs("test", summaries)
        d = result.to_dict()
        assert d["project"] == "test"
        assert "models" in d
        assert "segment_opportunities" in d
        assert "cost_drivers" in d

    def test_segment_opportunities_multiple_models_in_tag(self):
        summaries = (
            [_make_summary(f"a-{i}", model="gpt-4o", cost=0.02, tags=["simple"]) for i in range(5)]
            + [_make_summary(f"b-{i}", model="gpt-4o-mini", cost=0.002, tags=["simple"]) for i in range(5)]
        )
        evals = {
            **{f"a-{i}": _make_eval(f"a-{i}", judge_score=1) for i in range(5)},
            **{f"b-{i}": _make_eval(f"b-{i}", judge_score=1) for i in range(5)},
        }
        result = run_tradeoffs("test", summaries, evals)
        tag_opps = [o for o in result.segment_opportunities if o.segment_key == "tag:simple"]
        assert len(tag_opps) == 1
        assert len(tag_opps[0].models) == 2
