"""Outlier-focused + random baseline trace sampling for LLM context."""

from __future__ import annotations

import json
import random
from typing import Any

from .models import SampledTrace


def sample_traces(
    traces: list[dict[str, Any]],
    evaluations: dict[str, list[dict[str, Any]]] | None = None,
    max_total: int = 30,
) -> list[SampledTrace]:
    if not traces:
        return []

    evaluations = evaluations or {}

    # Pre-compute metrics for each trace
    scored: list[dict[str, Any]] = []
    for trace in traces:
        cost = sum(s.get("cost") or 0.0 for s in trace.get("spans", []))
        tokens = sum(
            (s.get("input_tokens") or 0) + (s.get("output_tokens") or 0)
            for s in trace.get("spans", [])
        )
        latency = _trace_latency(trace)
        has_error = any(s.get("status") == "error" for s in trace.get("spans", []))
        scored.append({
            "trace": trace,
            "cost": cost,
            "tokens": tokens,
            "latency": latency,
            "has_error": has_error,
        })

    seen: set[str] = set()
    result: list[SampledTrace] = []

    def _add(item: dict[str, Any], category: str) -> bool:
        tid = item["trace"]["trace_id"]
        if tid in seen or len(result) >= max_total:
            return False
        seen.add(tid)
        result.append(_to_sampled(item, category, evaluations))
        return True

    # Bucket 1: top 5 by cost
    by_cost = sorted(scored, key=lambda x: x["cost"], reverse=True)
    for item in by_cost[:5]:
        if item["cost"] > 0:
            _add(item, "high_cost")

    # Bucket 2: top 5 by latency
    by_latency = sorted(scored, key=lambda x: x["latency"] or 0, reverse=True)
    for item in by_latency[:5]:
        if (item["latency"] or 0) > 0:
            _add(item, "high_latency")

    # Bucket 3: errors (up to 10)
    errors = [s for s in scored if s["has_error"]]
    for item in errors[:10]:
        _add(item, "error")

    # Bucket 4: random fill
    remaining = [s for s in scored if s["trace"]["trace_id"] not in seen]
    if remaining:
        fill_count = max_total - len(result)
        sample_size = min(fill_count, len(remaining))
        if sample_size > 0:
            for item in random.sample(remaining, sample_size):
                _add(item, "random")

    return result


def _to_sampled(
    item: dict[str, Any],
    category: str,
    evaluations: dict[str, list[dict[str, Any]]],
) -> SampledTrace:
    trace = item["trace"]
    trace_id = trace["trace_id"]

    # Condense spans
    spans_summary = []
    for span in trace.get("spans", []):
        summary: dict[str, Any] = {
            "name": span.get("name", ""),
            "kind": span.get("kind", ""),
        }
        if span.get("cost") is not None:
            summary["cost"] = span["cost"]
        if span.get("input_tokens") is not None:
            summary["input_tokens"] = span["input_tokens"]
        if span.get("output_tokens") is not None:
            summary["output_tokens"] = span["output_tokens"]
        if span.get("model") is not None:
            summary["model"] = span["model"]
        latency = _span_latency(span)
        if latency is not None:
            summary["latency"] = round(latency, 3)
        if span.get("status") == "error":
            summary["error"] = span.get("error", "unknown error")
        spans_summary.append(summary)

    # Find first error
    error = None
    for span in trace.get("spans", []):
        if span.get("status") == "error":
            error = span.get("error", "unknown error")
            break

    return SampledTrace(
        trace_id=trace_id,
        category=category,
        cost=item["cost"],
        latency=item["latency"] or 0.0,
        tokens=item["tokens"],
        spans_summary=spans_summary,
        input_preview=_truncate(trace.get("input", ""), 200),
        output_preview=_truncate(trace.get("output", ""), 200),
        error=error,
    )


def _truncate(value: Any, max_len: int) -> str:
    if value is None:
        return ""
    if isinstance(value, dict):
        text = json.dumps(value)
    else:
        text = str(value)
    if len(text) > max_len:
        return text[:max_len] + "..."
    return text


def _trace_latency(trace: dict[str, Any]) -> float | None:
    from .aggregator import _parse_ts

    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()


def _span_latency(span: dict[str, Any]) -> float | None:
    from .aggregator import _parse_ts

    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()
