import json

import pytest

from bandito.cli.review.features import StructuralFeatures, extract_features
from bandito.cli.review.clusterer import cluster, TraceCluster, ClusteringResult
from bandito.cli.review.outliers import detect_structural_outliers, OutlierReport
from bandito.cli.review.clusterer import cluster_and_detect_outliers


def load_fixture(name: str) -> dict:
    with open(f"../shared/fixtures/{name}.json") as f:
        return json.load(f)


class TestStructuralFeatures:
    def test_simple_llm_trace(self):
        trace = load_fixture("simple-llm-trace")
        features = extract_features(trace)
        
        assert features.token_bucket == "tiny"
        assert features.span_composition == "llm_only"
        assert features.span_count_bucket == "single"
        assert features.has_error is False

    def test_rag_trace(self):
        trace = load_fixture("rag-trace")
        features = extract_features(trace)
        
        assert features.token_bucket == "small"
        assert features.span_composition == "retrieval_then_llm"
        assert features.span_count_bucket == "single"
        assert features.has_error is False

    def test_agent_loop_trace(self):
        trace = load_fixture("agent-loop-trace")
        features = extract_features(trace)
        
        assert features.token_bucket == "small"
        assert features.span_composition == "agent_loop"
        assert features.span_count_bucket == "few"
        assert features.has_error is False

    def test_error_trace(self):
        trace = {
            "trace_id": "test-err",
            "status": "error",
            "spans": [
                {"kind": "llm", "status": "error", "input_tokens": 10, "output_tokens": 5}
            ]
        }
        features = extract_features(trace)
        assert features.has_error is True

    def test_large_token_bucket(self):
        trace = {
            "trace_id": "test-large",
            "spans": [
                {"kind": "llm", "input_tokens": 8000, "output_tokens": 3000}
            ]
        }
        features = extract_features(trace)
        assert features.token_bucket == "large"

    def test_to_key_format(self):
        features = StructuralFeatures(
            token_bucket="medium",
            span_composition="retrieval_then_llm",
            span_count_bucket="few",
            has_error=False,
        )
        assert features.to_key() == "medium/retrieval_then_llm/few/ok"

    def test_summary_path(self):
        summary = {
            "trace_id": "test-summary",
            "total_tokens": 500,
            "span_count": 3,
            "span_kinds": ["retrieval", "llm"],
            "error_count": 0,
        }
        features = extract_features(summary)
        
        assert features.token_bucket == "small"
        assert features.span_composition == "retrieval_then_llm"
        assert features.span_count_bucket == "few"
        assert features.has_error is False

    def test_embedding_heavy_detection(self):
        trace = {
            "trace_id": "test-embedding",
            "spans": [
                {"kind": "embedding", "input_tokens": 1000},
                {"kind": "embedding", "input_tokens": 1000},
                {"kind": "embedding", "input_tokens": 1000},
                {"kind": "llm", "input_tokens": 100, "output_tokens": 50},
            ]
        }
        features = extract_features(trace)
        assert features.span_composition == "embedding_heavy"


class TestClusterer:
    def test_clusters_by_structural_features(self):
        traces = [
            load_fixture("simple-llm-trace"),
            {
                **load_fixture("simple-llm-trace"),
                "trace_id": "simple-2",
                "spans": [
                    {"kind": "llm", "input_tokens": 50, "output_tokens": 20}
                ]
            },
            load_fixture("rag-trace"),
        ]
        
        result = cluster(traces, min_cluster_size=1)
        
        assert result.total_traces == 3
        assert len(result.clusters) == 2
        
        cluster_ids = {c.cluster_id for c in result.clusters}
        assert "tiny/llm_only/single/ok" in cluster_ids
        assert "small/retrieval_then_llm/single/ok" in cluster_ids

    def test_min_cluster_size_filters_small(self):
        traces = [
            load_fixture("simple-llm-trace"),
            load_fixture("rag-trace"),
            {
                **load_fixture("simple-llm-trace"),
                "trace_id": "simple-2",
                "spans": [
                    {"kind": "llm", "input_tokens": 50, "output_tokens": 20}
                ]
            },
        ]
        
        result = cluster(traces, min_cluster_size=3)
        
        assert len(result.clusters) == 0

    def test_cluster_stats_computed(self):
        traces = [
            load_fixture("simple-llm-trace"),
            {
                **load_fixture("simple-llm-trace"),
                "trace_id": "simple-2",
                "spans": [
                    {"kind": "llm", "input_tokens": 50, "output_tokens": 20}
                ]
            },
            {
                **load_fixture("simple-llm-trace"),
                "trace_id": "simple-3",
                "spans": [
                    {"kind": "llm", "input_tokens": 30, "output_tokens": 10}
                ]
            },
        ]
        
        result = cluster(traces, min_cluster_size=3)
        
        assert len(result.clusters) == 1
        c = result.clusters[0]
        assert c.size == 3
        assert c.stats.trace_count == 3

    def test_empty_traces_returns_empty(self):
        result = cluster([])
        assert result.clusters == []
        assert result.total_traces == 0

    def test_to_dict_serializable(self):
        traces = [
            load_fixture("simple-llm-trace"),
            {
                **load_fixture("simple-llm-trace"),
                "trace_id": "simple-2",
                "spans": [
                    {"kind": "llm", "input_tokens": 50, "output_tokens": 20}
                ]
            },
            {
                **load_fixture("simple-llm-trace"),
                "trace_id": "simple-3",
                "spans": [
                    {"kind": "llm", "input_tokens": 30, "output_tokens": 10}
                ]
            },
        ]
        
        result = cluster(traces)
        d = result.to_dict()
        
        assert "clusters" in d
        assert d["total_traces"] == 3
        assert d["method"] == "rule_based"
        assert len(d["clusters"]) == 1


class TestOutliers:
    def test_no_outliers_in_homogeneous_population(self):
        traces = [
            {
                "trace_id": f"t{i}",
                "spans": [{"kind": "llm", "input_tokens": 100, "output_tokens": 50}]
            }
            for i in range(20)
        ]
        
        outliers = detect_structural_outliers(traces)
        
        assert isinstance(outliers, OutlierReport)
        assert outliers.total_traces == 20
        assert outliers.outlier_pct >= 0
        assert outliers.outlier_pct <= 5

    def test_outlier_detected_across_population(self):
        traces = [
            {
                "trace_id": f"t{i}",
                "spans": [{"kind": "llm", "input_tokens": 100, "output_tokens": 50}]
            }
            for i in range(15)
        ] + [
            {
                "trace_id": "outlier",
                "spans": [
                    {"kind": "llm", "input_tokens": 5000, "output_tokens": 2000}
                ]
            }
        ]
        
        outliers = detect_structural_outliers(traces)
        
        assert "outlier" in outliers.structural_outlier_ids

    def test_too_few_traces_returns_empty(self):
        traces = [
            {"trace_id": "t1", "spans": [{"kind": "llm", "input_tokens": 100}]},
            {"trace_id": "t2", "spans": [{"kind": "llm", "input_tokens": 200}]},
        ]
        
        outliers = detect_structural_outliers(traces, min_traces=10)
        
        assert outliers.structural_outlier_ids == []

    def test_outlier_percentage_calculation(self):
        traces = [
            {
                "trace_id": f"t{i}",
                "spans": [{"kind": "llm", "input_tokens": 100, "output_tokens": 50}]
            }
            for i in range(15)
        ] + [
            {
                "trace_id": "outlier",
                "spans": [
                    {"kind": "llm", "input_tokens": 5000, "output_tokens": 2000}
                ]
            }
        ]
        
        outliers = detect_structural_outliers(traces)
        
        assert outliers.total_traces == 16
        assert "outlier" in outliers.structural_outlier_ids

    def test_combined_function(self):
        traces = [
            {
                "trace_id": f"t{i}",
                "spans": [{"kind": "llm", "input_tokens": 100, "output_tokens": 50}]
            }
            for i in range(15)
        ] + [
            {
                "trace_id": "outlier",
                "spans": [
                    {"kind": "llm", "input_tokens": 5000, "output_tokens": 2000}
                ]
            }
        ]
        
        clustering, outliers = cluster_and_detect_outliers(traces, min_traces=10)
        
        assert isinstance(clustering, ClusteringResult)
        assert isinstance(outliers, OutlierReport)
        assert "outlier" in outliers.structural_outlier_ids
