"""Tests for auto-segmentation strategies and orchestrator."""

from __future__ import annotations

from bandito.review.segmenter import (
    _segment_by_model,
    _segment_by_span_kind_presence,
    _segment_by_tags,
    segment,
)


def _tool_span(name="search"):
    return {
        "span_id": f"sp-tool-{name}",
        "name": name,
        "kind": "tool",
        "started_at": "2026-03-25T10:00:00Z",
        "ended_at": "2026-03-25T10:00:01Z",
        "status": "ok",
    }


def _retrieval_span():
    return {
        "span_id": "sp-ret",
        "name": "retrieve",
        "kind": "retrieval",
        "started_at": "2026-03-25T10:00:00Z",
        "ended_at": "2026-03-25T10:00:01Z",
        "status": "ok",
    }


def _agent_span():
    return {
        "span_id": "sp-agent",
        "name": "agent-loop",
        "kind": "agent",
        "started_at": "2026-03-25T10:00:00Z",
        "ended_at": "2026-03-25T10:00:01Z",
        "status": "ok",
    }


class TestSegmentByTags:
    def test_no_tags(self, make_trace):
        traces = [make_trace() for _ in range(5)]
        result = _segment_by_tags(traces)
        assert result == {}

    def test_empty_tags_list(self, make_trace):
        traces = [make_trace(tags=[]) for _ in range(5)]
        result = _segment_by_tags(traces)
        assert result == {}

    def test_single_tag(self, make_trace):
        traces = [make_trace(tags=["prod"]) for _ in range(5)]
        result = _segment_by_tags(traces)
        assert "tag:prod" in result
        assert len(result["tag:prod"]) == 5

    def test_multiple_tags(self, make_trace):
        traces = [
            make_trace(tags=["prod"]),
            make_trace(tags=["prod"]),
            make_trace(tags=["staging"]),
        ]
        result = _segment_by_tags(traces)
        assert len(result["tag:prod"]) == 2
        assert len(result["tag:staging"]) == 1

    def test_multi_tag_trace(self, make_trace):
        t = make_trace(tags=["rag", "docs"])
        result = _segment_by_tags([t])
        assert t["trace_id"] in result["tag:rag"]
        assert t["trace_id"] in result["tag:docs"]


class TestSegmentByModel:
    def test_single_model(self, make_trace):
        traces = [make_trace(model="gpt-4o") for _ in range(5)]
        result = _segment_by_model(traces)
        assert "model:gpt-4o" in result
        assert len(result["model:gpt-4o"]) == 5

    def test_multiple_models(self, make_trace):
        traces = [
            make_trace(model="gpt-4o"),
            make_trace(model="gpt-4o"),
            make_trace(model="claude-sonnet"),
        ]
        result = _segment_by_model(traces)
        assert len(result["model:gpt-4o"]) == 2
        assert len(result["model:claude-sonnet"]) == 1

    def test_no_llm_spans(self, make_trace):
        traces = [make_trace(extra_spans=[_tool_span()], num_spans=0) for _ in range(3)]
        result = _segment_by_model(traces)
        assert result == {}

    def test_uses_first_llm_span(self, make_trace):
        """Should use model from first LLM span, not later ones."""
        t = make_trace(model="gpt-4o", extra_spans=[{
            "span_id": "sp-extra",
            "name": "second-call",
            "kind": "llm",
            "model": "gpt-4o-mini",
            "started_at": "2026-03-25T10:00:02Z",
            "ended_at": "2026-03-25T10:00:03Z",
            "status": "ok",
        }])
        result = _segment_by_model([t])
        assert "model:gpt-4o" in result
        assert "model:gpt-4o-mini" not in result


class TestSegmentBySpanKindPresence:
    def test_has_tool(self, make_trace):
        traces = [make_trace(extra_spans=[_tool_span()]) for _ in range(3)]
        result = _segment_by_span_kind_presence(traces)
        assert "has_tool" in result
        assert len(result["has_tool"]) == 3

    def test_has_retrieval(self, make_trace):
        traces = [make_trace(extra_spans=[_retrieval_span()]) for _ in range(3)]
        result = _segment_by_span_kind_presence(traces)
        assert "has_retrieval" in result

    def test_has_agent(self, make_trace):
        traces = [make_trace(extra_spans=[_agent_span()]) for _ in range(3)]
        result = _segment_by_span_kind_presence(traces)
        assert "has_agent" in result

    def test_no_special_kinds(self, make_trace):
        """Traces with only LLM spans produce no presence segments."""
        traces = [make_trace() for _ in range(5)]
        result = _segment_by_span_kind_presence(traces)
        assert result == {}

    def test_multiple_kinds(self, make_trace):
        """Trace with both tool and retrieval appears in both segments."""
        t = make_trace(extra_spans=[_tool_span(), _retrieval_span()])
        result = _segment_by_span_kind_presence([t])
        assert t["trace_id"] in result.get("has_tool", [])
        assert t["trace_id"] in result.get("has_retrieval", [])


class TestSegment:
    def test_minimum_size_filter(self, make_trace):
        """Segments below min_segment_size are dropped."""
        traces = [
            make_trace(tags=["prod"]),
            make_trace(tags=["prod"]),
            make_trace(tags=["staging"]),  # only 1 trace, should be dropped
        ]
        result = segment(traces, min_segment_size=2)
        keys = [s.key for s in result]
        assert "tag:prod" in keys
        assert "tag:staging" not in keys

    def test_stats_computed(self, make_trace):
        """Each segment has valid AggregateStats."""
        traces = [make_trace(model="gpt-4o", cost=0.01) for _ in range(5)]
        result = segment(traces, min_segment_size=3)
        model_seg = next((s for s in result if s.key == "model:gpt-4o"), None)
        assert model_seg is not None
        assert model_seg.stats.trace_count == 5
        assert model_seg.stats.cost.total > 0

    def test_empty_traces(self):
        result = segment([])
        assert result == []

    def test_too_few_traces(self, make_trace):
        traces = [make_trace()]
        result = segment(traces, min_segment_size=3)
        assert result == []

    def test_sorted_by_trace_count(self, make_trace):
        traces = (
            [make_trace(tags=["small"]) for _ in range(3)]
            + [make_trace(tags=["large"]) for _ in range(7)]
        )
        result = segment(traces, min_segment_size=3)
        tag_segs = [s for s in result if s.dimension == "tag"]
        if len(tag_segs) >= 2:
            assert tag_segs[0].trace_count >= tag_segs[1].trace_count

    def test_evaluations_passed_through(self, make_trace, make_evaluation):
        traces = [make_trace(trace_id=f"tr-{i}") for i in range(5)]
        evals = {
            f"tr-{i}": make_evaluation(f"tr-{i}", judge_score=1)
            for i in range(5)
        }
        result = segment(traces, evals, min_segment_size=3)
        model_seg = next((s for s in result if s.dimension == "model"), None)
        assert model_seg is not None
        assert model_seg.stats.judge_eval_count == 5

    def test_all_dimensions(self, make_trace):
        """Tags + models + span kinds all produce segments."""
        traces = [
            make_trace(tags=["prod"], model="gpt-4o", extra_spans=[_tool_span()])
            for _ in range(5)
        ]
        result = segment(traces, min_segment_size=3)
        dimensions = {s.dimension for s in result}
        assert "tag" in dimensions
        assert "model" in dimensions
        assert "span_kind_presence" in dimensions
