"""Auto-segmentation of traces by tags, models, and span kind presence.

Each strategy function returns dict[str, list[str]] mapping segment key → trace IDs.
The orchestrator runs all strategies, filters by minimum size, and calls aggregate()
per segment.

Future: ML-based clustering on user queries would add another strategy function
with the same return type.
"""

from __future__ import annotations

from typing import Any

from .aggregator import aggregate, aggregate_summaries
from .models import AggregateStats, Segment
from .utils import get_primary_model


def segment(
    traces: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]] | None = None,
    *,
    min_segment_size: int = 3,
    from_summaries: bool = False,
) -> list[Segment]:
    """Auto-segment traces and compute stats per segment.

    Returns segments sorted by trace count descending.
    Segments with fewer than min_segment_size traces are dropped.
    """
    evaluations = evaluations or {}

    if len(traces) < min_segment_size:
        return []

    # Build trace lookup
    trace_by_id: dict[str, dict[str, Any]] = {}
    for t in traces:
        tid = t.get("trace_id", "")
        if tid:
            trace_by_id[tid] = t

    # Run all strategies
    raw_segments: dict[str, tuple[str, str, list[str]]] = {}
    # Each entry: key -> (dimension, value, trace_ids)

    for key, trace_ids in _segment_by_tags(traces).items():
        raw_segments[key] = ("tag", key.removeprefix("tag:"), trace_ids)

    for key, trace_ids in _segment_by_model(traces).items():
        raw_segments[key] = ("model", key.removeprefix("model:"), trace_ids)

    for key, trace_ids in _segment_by_span_kind_presence(traces).items():
        raw_segments[key] = ("span_kind_presence", key, trace_ids)

    # Filter by min size, aggregate, build Segment objects
    segments: list[Segment] = []
    for key, (dimension, _value, trace_ids) in raw_segments.items():
        if len(trace_ids) < min_segment_size:
            continue

        seg_traces = [trace_by_id[tid] for tid in trace_ids if tid in trace_by_id]
        if len(seg_traces) < min_segment_size:
            continue

        seg_evals = {tid: evaluations[tid] for tid in trace_ids if tid in evaluations}
        agg = aggregate_summaries if from_summaries else aggregate
        stats = agg(seg_traces, seg_evals)

        segments.append(Segment(
            key=key,
            dimension=dimension,
            trace_count=len(seg_traces),
            trace_ids=trace_ids,
            stats=stats,
        ))

    segments.sort(key=lambda s: s.trace_count, reverse=True)
    return segments


def _segment_by_tags(traces: list[dict[str, Any]]) -> dict[str, list[str]]:
    """One segment per unique tag. Traces with multiple tags appear in multiple segments."""
    tag_map: dict[str, list[str]] = {}
    for t in traces:
        tid = t.get("trace_id", "")
        if not tid:
            continue
        for tag in t.get("tags", []):
            if not tag:
                continue
            key = f"tag:{tag}"
            tag_map.setdefault(key, []).append(tid)
    return tag_map


def _segment_by_model(traces: list[dict[str, Any]]) -> dict[str, list[str]]:
    """One segment per primary model. Works with full traces or summaries."""
    model_map: dict[str, list[str]] = {}
    for t in traces:
        tid = t.get("trace_id", "")
        if not tid:
            continue
        # Summary path
        model = t.get("primary_model")
        # Full trace path
        if not model and "spans" in t:
            model = get_primary_model(t)
        if model:
            model_map.setdefault(f"model:{model}", []).append(tid)
    return model_map


def _segment_by_span_kind_presence(
    traces: list[dict[str, Any]],
) -> dict[str, list[str]]:
    """Binary segments for presence of retrieval, tool, and agent span kinds.

    Works with full traces (reads spans[].kind) or summaries (reads span_kinds[]).
    """
    _INTERESTING_KINDS = {"retrieval", "tool", "agent"}
    kind_map: dict[str, list[str]] = {}
    for t in traces:
        tid = t.get("trace_id", "")
        if not tid:
            continue
        # Summary path: span_kinds is a list of unique kinds
        if "span_kinds" in t and "spans" not in t:
            kinds_present = set(t["span_kinds"]) & _INTERESTING_KINDS
        else:
            # Full trace path: extract from spans
            kinds_present = set()
            for span in t.get("spans", []):
                kind = span.get("kind", "")
                if kind in _INTERESTING_KINDS:
                    kinds_present.add(kind)
        for kind in kinds_present:
            key = f"has_{kind}"
            kind_map.setdefault(key, []).append(tid)
    return kind_map
