"""Extract structural features from traces.

Structural features capture how a trace behaves (tokens, spans, errors),
not what it says. Used for clustering and outlier detection.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any


@dataclass
class StructuralFeatures:
    token_bucket: str
    span_composition: str
    span_count_bucket: str
    has_error: bool

    def to_key(self) -> str:
        return f"{self.token_bucket}/{self.span_composition}/{self.span_count_bucket}/{'error' if self.has_error else 'ok'}"

    @classmethod
    def from_trace(cls, trace: dict[str, Any]) -> StructuralFeatures:
        spans = trace.get("spans", [])
        
        total_tokens = cls._total_tokens(trace, spans)
        token_bucket = cls._token_bucket(total_tokens)
        
        span_composition = cls._span_composition(spans)
        span_count_bucket = cls._span_count_bucket(len(spans))
        has_error = cls._has_error(trace, spans)
        
        return cls(
            token_bucket=token_bucket,
            span_composition=span_composition,
            span_count_bucket=span_count_bucket,
            has_error=has_error,
        )

    @classmethod
    def from_summary(cls, summary: dict[str, Any]) -> StructuralFeatures:
        total_tokens = summary.get("total_tokens", 0)
        token_bucket = cls._token_bucket(total_tokens)
        
        span_composition = cls._span_composition_from_summary(summary)
        span_count_bucket = cls._span_count_bucket(summary.get("span_count", 0))
        has_error = summary.get("error_count", 0) > 0
        
        return cls(
            token_bucket=token_bucket,
            span_composition=span_composition,
            span_count_bucket=span_count_bucket,
            has_error=has_error,
        )

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> StructuralFeatures:
        return cls(
            token_bucket=data.get("token_bucket", "unknown"),
            span_composition=data.get("span_composition", "unknown"),
            span_count_bucket=data.get("span_count_bucket", "unknown"),
            has_error=data.get("has_error", False),
        )

    def to_dict(self) -> dict[str, Any]:
        return {
            "token_bucket": self.token_bucket,
            "span_composition": self.span_composition,
            "span_count_bucket": self.span_count_bucket,
            "has_error": self.has_error,
        }

    @staticmethod
    def _total_tokens(trace: dict[str, Any], spans: list[dict]) -> int:
        if spans:
            total = 0
            for span in spans:
                inp = span.get("input_tokens") or 0
                out = span.get("output_tokens") or 0
                total += inp + out
            return total
        
        return trace.get("total_tokens", 0)

    @staticmethod
    def _token_bucket(tokens: int) -> str:
        if tokens < 100:
            return "tiny"
        elif tokens < 1000:
            return "small"
        elif tokens < 10000:
            return "medium"
        else:
            return "large"

    @staticmethod
    def _span_composition(spans: list[dict[str, Any]]) -> str:
        if not spans:
            return "llm_only"
        
        kinds = [s.get("kind", "") for s in spans]
        
        has_retrieval = "retrieval" in kinds
        has_tool = "tool" in kinds
        has_embedding = "embedding" in kinds
        
        llm_count = kinds.count("llm")
        has_agent = llm_count > 1
        
        if has_retrieval:
            retrieval_idx = kinds.index("retrieval")
            llm_idxs = [i for i, k in enumerate(kinds) if k == "llm"]
            if llm_idxs and all(llm_idx > retrieval_idx for llm_idx in llm_idxs):
                return "retrieval_then_llm"
        
        if has_embedding and kinds.count("embedding") > len(kinds) / 2:
            return "embedding_heavy"
        
        if has_agent or (has_tool and llm_count > 1):
            return "agent_loop"
        
        return "llm_only"

    @staticmethod
    def _span_composition_from_summary(summary: dict[str, Any]) -> str:
        span_kinds = set(summary.get("span_kinds", []))
        
        if "retrieval" in span_kinds:
            return "retrieval_then_llm"
        
        if "embedding" in span_kinds and summary.get("span_count", 0) > 0:
            embedding_count = summary.get("span_kinds", []).count("embedding")
            if embedding_count > summary.get("span_count", 1) / 2:
                return "embedding_heavy"
        
        llm_count = summary.get("span_kinds", []).count("llm")
        if llm_count > 1:
            return "agent_loop"
        
        return "llm_only"

    @staticmethod
    def _span_count_bucket(count: int) -> str:
        if count <= 2:
            return "single"
        elif count <= 5:
            return "few"
        elif count <= 10:
            return "many"
        else:
            return "loop"

    @staticmethod
    def _has_error(trace: dict[str, Any], spans: list[dict[str, Any]]) -> bool:
        if trace.get("status") == "error":
            return True
        
        for span in spans:
            if span.get("status") == "error":
                return True
        
        return trace.get("error_count", 0) > 0


def extract_features(trace: dict[str, Any]) -> StructuralFeatures:
    """Extract structural features from a trace or summary.
    
    Works with both full traces (with spans[]) and summaries (with span_kinds).
    """
    if "spans" in trace:
        return StructuralFeatures.from_trace(trace)
    else:
        return StructuralFeatures.from_summary(trace)
