"""Tests for summary-path aggregation and segmentation."""

from __future__ import annotations

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

from bandito.review.aggregator import aggregate, aggregate_summaries
from bandito.review.segmenter import segment
from bandito.storage.summary import summarize_trace


def _make_full_trace(
    trace_id: str,
    model: str = "gpt-4o",
    cost: float | None = 0.01,
    tags: list[str] | None = None,
) -> dict[str, Any]:
    now = datetime.now(timezone.utc)
    start = (now - timedelta(seconds=1)).isoformat().replace("+00:00", "Z")
    end = now.isoformat().replace("+00:00", "Z")
    span: dict[str, Any] = {
        "span_id": f"sp-{trace_id}",
        "name": "llm", "kind": "llm",
        "started_at": start, "ended_at": end,
        "status": "ok", "model": model,
        "input_tokens": 100, "output_tokens": 50,
    }
    if cost is not None:
        span["cost"] = cost
    trace: dict[str, Any] = {
        "trace_id": trace_id, "project": "test",
        "started_at": start, "ended_at": end,
        "status": "ok", "spans": [span],
        "input": "hello", "output": "world",
    }
    if tags:
        trace["tags"] = tags
    return trace


def _make_summary(trace_id: str, model: str = "gpt-4o", cost: float | None = 0.01, tags: list[str] | None = None) -> dict[str, Any]:
    return summarize_trace(_make_full_trace(trace_id, model=model, cost=cost, tags=tags))


class TestSummarizeTrace:
    def test_extracts_fields(self):
        trace = _make_full_trace("t-1", model="gpt-4o", cost=0.02)
        s = summarize_trace(trace)
        assert s["trace_id"] == "t-1"
        assert s["primary_model"] == "gpt-4o"
        assert s["total_cost"] == 0.02
        assert s["total_tokens"] == 150
        assert s["span_count"] == 1
        assert "llm" in s["span_kinds"]
        assert s["input_preview"] == "hello"
        assert s["output_preview"] == "world"
        assert "spans" not in s

    def test_no_cost_is_none(self):
        trace = _make_full_trace("t-1", cost=None)
        s = summarize_trace(trace)
        assert s["total_cost"] is None

    def test_tags_preserved(self):
        trace = _make_full_trace("t-1", tags=["prod", "v2"])
        s = summarize_trace(trace)
        assert s["tags"] == ["prod", "v2"]


class TestAggregateSummaries:
    def test_basic_stats(self):
        summaries = [_make_summary(f"t-{i}", cost=0.01) for i in range(5)]
        stats = aggregate_summaries(summaries)
        assert stats.trace_count == 5
        assert stats.cost.total == pytest.approx(0.05, abs=0.001)
        assert stats.cost.count == 5

    def test_matches_full_trace_totals(self):
        """Summary aggregation should match full trace aggregation for key totals."""
        traces = [_make_full_trace(f"t-{i}", cost=0.01 * (i + 1)) for i in range(5)]
        summaries = [summarize_trace(t) for t in traces]

        full_stats = aggregate(traces)
        summary_stats = aggregate_summaries(summaries)

        assert full_stats.trace_count == summary_stats.trace_count
        assert full_stats.cost.total == pytest.approx(summary_stats.cost.total, abs=0.001)
        assert full_stats.cost.count == summary_stats.cost.count

    def test_no_cost_excluded(self):
        summaries = [_make_summary(f"t-{i}", cost=None) for i in range(5)]
        stats = aggregate_summaries(summaries)
        assert stats.cost.count == 0

    def test_evaluations_counted(self):
        summaries = [_make_summary(f"t-{i}") for i in range(5)]
        evals = {f"t-{i}": {"judge_score": 1} for i in range(3)}
        stats = aggregate_summaries(summaries, evals)
        assert stats.judge_eval_count == 3
        assert stats.judge_avg_grade == 1.0

    def test_empty_input(self):
        stats = aggregate_summaries([])
        assert stats.trace_count == 0


class TestSegmentSummaries:
    def test_segments_by_tag(self):
        summaries = (
            [_make_summary(f"a-{i}", tags=["prod"]) for i in range(5)]
            + [_make_summary(f"b-{i}", tags=["staging"]) for i in range(3)]
        )
        segs = segment(summaries, from_summaries=True, min_segment_size=3)
        keys = {s.key for s in segs}
        assert "tag:prod" in keys
        assert "tag:staging" in keys

    def test_segments_by_model(self):
        summaries = (
            [_make_summary(f"a-{i}", model="gpt-4o") for i in range(5)]
            + [_make_summary(f"b-{i}", model="gpt-4o-mini") for i in range(5)]
        )
        segs = segment(summaries, from_summaries=True, min_segment_size=3)
        keys = {s.key for s in segs}
        assert "model:gpt-4o" in keys
        assert "model:gpt-4o-mini" in keys

    def test_segment_stats_use_summary_aggregation(self):
        summaries = [_make_summary(f"t-{i}", cost=0.01) for i in range(5)]
        segs = segment(summaries, from_summaries=True, min_segment_size=3)
        for seg in segs:
            assert seg.stats.trace_count > 0


import pytest
