"""Deterministic trace aggregation — no LLM, pure stats."""

from __future__ import annotations

from datetime import datetime
from typing import Any

from .models import AggregateStats, Distribution, ModelBreakdown, SpanKindBreakdown


def aggregate(
    traces: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]] | None = None,
) -> AggregateStats:
    """Aggregate full traces (with spans[]) into stats."""
    evaluations = evaluations or {}
    if not traces:
        return AggregateStats()
    return _aggregate_from_traces(traces, evaluations)


def aggregate_summaries(
    summaries: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]] | None = None,
) -> AggregateStats:
    """Aggregate trace summaries (pre-computed stats, no spans) into stats.

    Note: summaries lack per-span detail. Model breakdowns will not include
    per-model cost or latency — only trace counts per model. Span kind
    breakdowns are approximate (traces-with-kind, not span counts).
    """
    evaluations = evaluations or {}
    if not summaries:
        return AggregateStats()
    return _aggregate_from_summaries(summaries, evaluations)


def _aggregate_from_summaries(
    summaries: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]],
) -> AggregateStats:
    """Aggregate from pre-computed trace summaries."""
    costs: list[float] = []
    latencies: list[float] = []
    token_counts: list[int] = []
    model_map: dict[str, int] = {}  # model -> trace count
    kind_set: dict[str, int] = {}  # kind -> trace count (approximate)
    error_count = 0
    eval_count = 0
    human_grades: list[float] = []
    judge_grades: list[float] = []

    for s in summaries:
        cost = s.get("total_cost")
        if cost is not None:
            costs.append(cost)

        latency = s.get("latency")
        if latency is not None:
            latencies.append(latency)

        tokens = s.get("total_tokens")
        if tokens is not None:
            token_counts.append(tokens)

        error_count += s.get("error_count", 0)

        model = s.get("primary_model")
        if model:
            model_map[model] = model_map.get(model, 0) + 1

        for kind in s.get("span_kinds", []):
            kind_set[kind] = kind_set.get(kind, 0) + 1

        trace_id = s.get("trace_id", "")
        ev = evaluations.get(trace_id)
        if ev:
            has_human = ev.get("human_score") is not None
            has_judge = ev.get("judge_score") is not None
            if has_human or has_judge:
                eval_count += 1
            if has_human:
                human_grades.append(ev["human_score"])
            if has_judge:
                judge_grades.append(ev["judge_score"])

    # Model breakdowns (limited — no per-model cost/latency from summaries)
    models = [
        ModelBreakdown(
            model=model,
            provider=None,
            count=count,
            has_cost=False,
        )
        for model, count in sorted(model_map.items())
    ]

    # Span kind breakdowns (approximate — counts are traces-with-kind, not span counts)
    span_kinds = [
        SpanKindBreakdown(kind=kind, count=count)
        for kind, count in sorted(kind_set.items())
    ]

    return AggregateStats(
        trace_count=len(summaries),
        span_count=sum(s.get("span_count", 0) for s in summaries),
        cost=_compute_distribution(costs),
        latency=_compute_distribution(latencies),
        tokens=_compute_distribution([float(t) for t in token_counts]),
        models=models,
        span_kinds=span_kinds,
        error_count=error_count,
        eval_count=eval_count,
        human_eval_count=len(human_grades),
        human_avg_grade=sum(human_grades) / len(human_grades) if human_grades else None,
        judge_eval_count=len(judge_grades),
        judge_avg_grade=sum(judge_grades) / len(judge_grades) if judge_grades else None,
    )


def _aggregate_from_traces(
    traces: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]],
) -> AggregateStats:
    """Aggregate from full traces with spans."""
    costs: list[float] = []
    latencies: list[float] = []
    token_counts: list[int] = []
    model_map: dict[tuple[str, str | None], dict[str, Any]] = {}
    kind_map: dict[str, dict[str, Any]] = {}
    total_spans = 0
    error_count = 0
    eval_count = 0
    human_grades: list[float] = []
    judge_grades: list[float] = []

    for trace in traces:
        trace_cost = 0.0
        trace_tokens = 0
        trace_has_cost = False
        trace_has_tokens = False

        for span in trace.get("spans", []):
            total_spans += 1
            kind = span.get("kind", "unknown")

            if kind not in kind_map:
                kind_map[kind] = {"count": 0, "total_latency": 0.0, "error_count": 0}
            kind_map[kind]["count"] += 1

            span_latency = _span_latency(span)
            if span_latency is not None:
                kind_map[kind]["total_latency"] += span_latency

            if span.get("status") == "error":
                error_count += 1
                kind_map[kind]["error_count"] += 1

            cost = span.get("cost")
            if cost is not None:
                trace_cost += cost
                trace_has_cost = True

            model = span.get("model")
            if model is not None:
                key = (model, span.get("provider"))
                if key not in model_map:
                    model_map[key] = {
                        "count": 0,
                        "total_cost": 0.0,
                        "total_latency": 0.0,
                        "total_tokens": 0,
                        "has_cost": False,
                    }
                model_map[key]["count"] += 1
                if cost is not None:
                    model_map[key]["total_cost"] += cost
                    model_map[key]["has_cost"] = True
                if span_latency is not None:
                    model_map[key]["total_latency"] += span_latency
                inp = span.get("input_tokens") or 0
                out = span.get("output_tokens") or 0
                model_map[key]["total_tokens"] += inp + out

            inp = span.get("input_tokens")
            out = span.get("output_tokens")
            if inp is not None or out is not None:
                trace_has_tokens = True
                trace_tokens += (inp or 0) + (out or 0)

        # Only include in distribution if the trace actually reported the metric
        if trace_has_cost:
            costs.append(trace_cost)
        if trace_has_tokens:
            token_counts.append(trace_tokens)

        trace_latency = _trace_latency(trace)
        if trace_latency is not None:
            latencies.append(trace_latency)

        trace_id = trace.get("trace_id", "")
        ev = evaluations.get(trace_id)
        if ev:
            has_human = ev.get("human_score") is not None
            has_judge = ev.get("judge_score") is not None
            if has_human or has_judge:
                eval_count += 1
            if has_human:
                human_grades.append(ev["human_score"])
            if has_judge:
                judge_grades.append(ev["judge_score"])

    # Span kind breakdowns
    total_kind_latency = sum(v["total_latency"] for v in kind_map.values())
    span_kinds = []
    for kind, data in sorted(kind_map.items()):
        count = data["count"]
        tl = data["total_latency"]
        span_kinds.append(
            SpanKindBreakdown(
                kind=kind,
                count=count,
                total_latency=tl,
                avg_latency=tl / count if count else 0.0,
                pct_of_total=tl / total_kind_latency * 100 if total_kind_latency else 0.0,
                error_count=data["error_count"],
            )
        )

    # Model breakdowns
    models = []
    for (model, provider), data in sorted(model_map.items()):
        count = data["count"]
        has_cost = data["has_cost"]
        models.append(
            ModelBreakdown(
                model=model,
                provider=provider,
                count=count,
                total_cost=data["total_cost"] if has_cost else 0.0,
                avg_cost=data["total_cost"] / count if has_cost and count else 0.0,
                avg_latency=data["total_latency"] / count if count else 0.0,
                total_tokens=data["total_tokens"],
                has_cost=has_cost,
            )
        )

    stats = AggregateStats(
        trace_count=len(traces),
        span_count=total_spans,
        cost=_compute_distribution(costs),
        latency=_compute_distribution(latencies),
        tokens=_compute_distribution([float(t) for t in token_counts]),
        models=models,
        span_kinds=span_kinds,
        error_count=error_count,
        eval_count=eval_count,
        human_eval_count=len(human_grades),
        human_avg_grade=sum(human_grades) / len(human_grades) if human_grades else None,
        judge_eval_count=len(judge_grades),
        judge_avg_grade=sum(judge_grades) / len(judge_grades) if judge_grades else None,
    )
    return stats


def _compute_distribution(values: list[float]) -> Distribution:
    if not values:
        return Distribution()
    s = sorted(values)
    n = len(s)
    return Distribution(
        min=s[0],
        max=s[-1],
        mean=sum(s) / n,
        median=s[n // 2] if n % 2 else (s[n // 2 - 1] + s[n // 2]) / 2,
        p95=s[min(int(n * 0.95), n - 1)],
        p99=s[min(int(n * 0.99), n - 1)],
        total=sum(s),
        count=n,
    )


def _parse_ts(raw: str) -> datetime | None:
    try:
        return datetime.fromisoformat(raw.replace("Z", "+00:00"))
    except (ValueError, TypeError, AttributeError):
        return None


def _span_latency(span: dict[str, Any]) -> float | None:
    s = span.get("started_at")
    e = span.get("ended_at")
    if not s or not e:
        return None
    start = _parse_ts(s)
    end = _parse_ts(e)
    if start is None or end is None:
        return None
    return (end - start).total_seconds()


def _trace_latency(trace: dict[str, Any]) -> float | None:
    s = trace.get("started_at")
    e = trace.get("ended_at")
    if not s or not e:
        return None
    start = _parse_ts(s)
    end = _parse_ts(e)
    if start is None or end is None:
        return None
    return (end - start).total_seconds()
