"""Structural clustering of traces by behavioral features.

Groups traces by token bucket, span composition, span count, and error status.
No ML dependencies — pure rule-based clustering.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .aggregator import aggregate, aggregate_summaries
from .features import StructuralFeatures, extract_features
from .models import AggregateStats
from .outliers import OutlierReport, detect_structural_outliers


@dataclass
class TraceCluster:
    cluster_id: str
    structural_key: StructuralFeatures
    name: str | None = None
    size: int = 0
    trace_ids: list[str] = field(default_factory=list)
    stats: AggregateStats = field(default_factory=AggregateStats)

    def to_dict(self) -> dict[str, Any]:
        return {
            "cluster_id": self.cluster_id,
            "structural_key": self.structural_key.to_dict(),
            "name": self.name,
            "size": self.size,
            "trace_ids": self.trace_ids,
            "stats": {
                "trace_count": self.stats.trace_count,
                "span_count": self.stats.span_count,
                "cost": {
                    "total": self.stats.cost.total,
                    "mean": self.stats.cost.mean,
                    "median": self.stats.cost.median,
                    "p95": self.stats.cost.p95,
                    "min": self.stats.cost.min,
                    "max": self.stats.cost.max,
                },
                "latency": {
                    "mean": self.stats.latency.mean,
                    "median": self.stats.latency.median,
                    "p95": self.stats.latency.p95,
                },
                "tokens": {
                    "mean": self.stats.tokens.mean,
                    "median": self.stats.tokens.median,
                },
                "models": [
                    {"model": m.model, "count": m.count, "avg_cost": m.avg_cost}
                    for m in self.stats.models
                ],
                "span_kinds": [
                    {"kind": k.kind, "count": k.count, "pct_of_total": k.pct_of_total}
                    for k in self.stats.span_kinds
                ],
                "eval_count": self.stats.eval_count,
                "human_avg_grade": self.stats.human_avg_grade,
                "judge_avg_grade": self.stats.judge_avg_grade,
            },
        }


@dataclass
class ClusteringResult:
    clusters: list[TraceCluster]
    total_traces: int
    method: str = "rule_based"

    def to_dict(self) -> dict[str, Any]:
        return {
            "clusters": [c.to_dict() for c in self.clusters],
            "total_traces": self.total_traces,
            "n_clusters": len(self.clusters),
            "method": self.method,
        }


def cluster(
    traces: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]] | None = None,
    *,
    min_cluster_size: int = 3,
    from_summaries: bool = False,
) -> ClusteringResult:
    """Cluster traces by structural features.

    Returns clusters sorted by size descending.
    Clusters smaller than min_cluster_size are dropped.
    """
    evaluations = evaluations or {}

    if len(traces) < min_cluster_size:
        return ClusteringResult(clusters=[], total_traces=len(traces))

    trace_by_id: dict[str, dict[str, Any]] = {}
    features_by_id: dict[str, StructuralFeatures] = {}
    clusters_map: dict[str, list[str]] = {}

    for trace in traces:
        tid = trace.get("trace_id", "")
        if not tid:
            continue
        
        trace_by_id[tid] = trace
        features = extract_features(trace)
        features_by_id[tid] = features
        
        key = features.to_key()
        clusters_map.setdefault(key, []).append(tid)

    cluster_list: list[TraceCluster] = []
    for key, trace_ids in clusters_map.items():
        if len(trace_ids) < min_cluster_size:
            continue
        
        cluster_traces = [trace_by_id[tid] for tid in trace_ids if tid in trace_by_id]
        if len(cluster_traces) < min_cluster_size:
            continue

        sample_features = features_by_id.get(trace_ids[0])
        if not sample_features:
            continue

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

        cluster_list.append(TraceCluster(
            cluster_id=key,
            structural_key=sample_features,
            size=len(cluster_traces),
            trace_ids=trace_ids,
            stats=stats,
        ))

    cluster_list.sort(key=lambda c: c.size, reverse=True)
    return ClusteringResult(clusters=cluster_list, total_traces=len(traces))


def cluster_and_detect_outliers(
    traces: list[dict[str, Any]],
    evaluations: dict[str, dict[str, Any]] | None = None,
    *,
    min_traces: int = 10,
    min_cluster_size: int = 3,
    percentile_threshold: float = 95.0,
    from_summaries: bool = False,
) -> tuple[ClusteringResult, OutlierReport]:
    """Run clustering and outlier detection with aligned thresholds.

    Clustering drops clusters smaller than min_cluster_size.
    Outlier detection requires at least min_traces to produce results.

    Both operations run on the full trace set.

    Args:
        traces: Trace list
        evaluations: Optional evaluation dict
        min_traces: Minimum traces needed for meaningful analysis (default 10)
        min_cluster_size: Minimum traces per cluster (default 3)
        percentile_threshold: Percentile for outlier threshold (default 95)
        from_summaries: Use summary aggregation path

    Returns:
        Tuple of (ClusteringResult, OutlierReport)
    """
    clustering = cluster(
        traces,
        evaluations,
        min_cluster_size=min_cluster_size,
        from_summaries=from_summaries,
    )
    outliers = detect_structural_outliers(
        traces,
        percentile_threshold=percentile_threshold,
        min_traces=min_traces,
    )
    return clustering, outliers
