#!/usr/bin/env python3
"""Comprehensive data quality validation for router training data.

This module provides tools to detect and report data quality issues in the
router dataset, including:
- Response truncation and corruption
- API failures and error responses
- Evaluation metric anomalies
- Cross-model consistency checks
- Benchmark-specific validation

Usage:
    # Run full validation
    python -m embergraphrouter.data_processing.data_quality

    # Quick validation (skip expensive checks)
    python -m embergraphrouter.data_processing.data_quality --quick

    # Validate specific LLMs
    python -m embergraphrouter.data_processing.data_quality --llms "GPT-5.1" "Gemini 3 Pro Preview"

    # Export report
    python -m embergraphrouter.data_processing.data_quality --output report.json
"""

from __future__ import annotations

import hashlib
import json
import logging
import re
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import pandas as pd

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

_MODULE_DIR = Path(__file__).resolve().parent.parent


# =============================================================================
# Public Benchmark Ranges (2025)
# =============================================================================
# Expected accuracy ranges for best-performing LLM on each task.
# Format: task_id -> (min_expected, max_expected)
# Based on public benchmark data from model providers.
PUBLIC_BENCHMARK_RANGES: dict[str, tuple[float, float]] = {
    # Math reasoning
    "GSM8K": (0.85, 0.99),  # Frontier models get 95%+
    # Multiple choice knowledge
    "mmlu": (0.85, 0.99),  # MMLU ~90%
    "arc": (0.85, 0.99),  # ARC-Challenge ~95%
    "hellaswag": (0.90, 1.00),  # Near ceiling for frontier
    "truthfulqa": (0.70, 0.95),  # Harder, ~80%
    "winogrande": (0.85, 1.00),  # Near ceiling
    "agieval": (0.75, 0.99),  # SAT/LSAT/GRE level
    # MMLU-Pro subjects (harder than MMLU)
    "mmlu_pro_biology": (0.70, 0.95),
    "mmlu_pro_business": (0.70, 0.95),
    "mmlu_pro_chemistry": (0.60, 0.90),  # Hardest subject
    "mmlu_pro_history": (0.70, 0.95),
    "mmlu_pro_law": (0.65, 0.90),
    "mmlu_pro_other": (0.70, 0.95),
    "mmlu_pro_psychology": (0.75, 0.95),
    # BBEH reasoning tasks (challenging)
    "bbeh_boardgame_qa": (0.30, 0.85),
    "bbeh_boolean_expressions": (0.40, 0.90),
    "bbeh_buggy_tables": (0.30, 0.80),
    "bbeh_causal_understanding": (0.40, 0.85),
    "bbeh_disambiguation_qa": (0.40, 0.85),
    "bbeh_dyck_languages": (0.30, 0.80),
    "bbeh_geometric_shapes": (0.30, 0.80),
    "bbeh_hyperbaton": (0.40, 0.85),
    "bbeh_linguini": (0.30, 0.80),
    "bbeh_movie_recommendation": (0.40, 0.85),
    "bbeh_multistep_arithmetic": (0.30, 0.85),
    # Free-form tasks (F1 scores, typically lower)
    "alpaca_data": (0.20, 0.60),
    "multi_news": (0.20, 0.50),
    "SQUAD": (0.60, 0.95),
}

# LLM capability tiers (for relative ordering checks)
# Tier 1: Frontier models (should be top performers)
# Tier 2: Strong mid-tier models
# Tier 3: Efficient/fast models (lowest accuracy expected)
LLM_CAPABILITY_TIERS: dict[str, int] = {
    "Claude Opus 4.5": 1,
    "GPT-5.1": 1,
    "GPT-5.1 Codex": 1,
    "Gemini 3 Pro Preview": 1,
    "Claude Haiku 4.5": 2,
    "Gemini 3 Flash Preview": 2,  # Fast but capable Gemini 3 variant
    "Gemini 2.5 Flash Lite": 3,
    "GPT-5 Nano": 3,
}

# Tasks disabled due to requiring code execution (zeros expected)
DISABLED_TASKS = frozenset({"livecodebench", "mbpp"})


def _query_hash(query: str, llm_name: str) -> str:
    """Generate cache key for query+llm pair."""
    return hashlib.sha256(f"{query}||{llm_name}".encode()).hexdigest()[:16]


@dataclass
class ResponseIssue:
    """A single response quality issue."""

    issue_type: str
    severity: str  # "critical", "warning", "info"
    description: str
    query_preview: str = ""
    response_preview: str = ""
    metadata: dict = field(default_factory=dict)


@dataclass
class LLMQualityReport:
    """Quality report for a single LLM."""

    llm_name: str
    total_responses: int = 0
    issues: list[ResponseIssue] = field(default_factory=list)
    issue_counts: dict[str, int] = field(default_factory=lambda: defaultdict(int))
    metrics: dict[str, float] = field(default_factory=dict)

    def add_issue(self, issue: ResponseIssue) -> None:
        """Record an issue."""
        self.issues.append(issue)
        self.issue_counts[issue.issue_type] += 1

    def get_issue_rate(self, issue_type: str) -> float:
        """Get percentage of responses with this issue."""
        if self.total_responses == 0:
            return 0.0
        return self.issue_counts[issue_type] / self.total_responses * 100

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON export."""
        return {
            "llm_name": self.llm_name,
            "total_responses": self.total_responses,
            "issue_counts": dict(self.issue_counts),
            "issue_rates": {
                k: self.get_issue_rate(k) for k in self.issue_counts
            },
            "metrics": self.metrics,
            "sample_issues": [
                {
                    "type": i.issue_type,
                    "severity": i.severity,
                    "description": i.description,
                    "query": i.query_preview,
                    "response": i.response_preview,
                }
                for i in self.issues[:10]  # Limit samples
            ],
        }


@dataclass
class DataQualityReport:
    """Complete data quality report."""

    llm_reports: dict[str, LLMQualityReport] = field(default_factory=dict)
    global_issues: list[ResponseIssue] = field(default_factory=list)
    summary: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON export."""
        return {
            "summary": self.summary,
            "llm_reports": {k: v.to_dict() for k, v in self.llm_reports.items()},
            "global_issues": [
                {
                    "type": i.issue_type,
                    "severity": i.severity,
                    "description": i.description,
                }
                for i in self.global_issues
            ],
        }


class ResponseQualityChecker:
    """Checks individual response quality."""

    # Thresholds
    VERY_SHORT_THRESHOLD = 50
    EXTREMELY_SHORT_THRESHOLD = 20
    SUSPICIOUSLY_FAST_LATENCY = 0.3
    VERY_SLOW_LATENCY = 120

    # Error patterns
    ERROR_PATTERNS = [
        r"error",
        r"exception",
        r"timeout",
        r"rate.?limit",
        r"quota.?exceeded",
        r"429",
        r"500",
        r"503",
        r"internal.?server",
        r"service.?unavailable",
        r"try.?again",
        r"temporarily.?unavailable",
    ]

    # Truncation patterns
    VALID_ENDINGS = ".!?>\n0123456789\"')]}。！？"

    def __init__(self):
        self._error_regex = re.compile(
            "|".join(self.ERROR_PATTERNS), re.IGNORECASE
        )

    def check_response(
        self,
        response: str,
        latency: float,
        query: str = "",
        task_id: str = "",
    ) -> list[ResponseIssue]:
        """Check a single response for quality issues.

        Args:
            response: The LLM response text.
            latency: Response latency in seconds.
            query: Original query (for context).
            task_id: Task identifier for benchmark-specific checks.

        Returns:
            List of detected issues.
        """
        issues = []
        query_preview = query[:100] + "..." if len(query) > 100 else query
        resp_preview = response[:100] + "..." if len(response) > 100 else response

        # Empty response
        if not response or len(response.strip()) == 0:
            issues.append(ResponseIssue(
                issue_type="empty_response",
                severity="critical",
                description="Response is empty or whitespace only",
                query_preview=query_preview,
                response_preview=resp_preview,
            ))
            return issues  # No further checks needed

        # Extremely short
        if len(response) < self.EXTREMELY_SHORT_THRESHOLD:
            issues.append(ResponseIssue(
                issue_type="extremely_short",
                severity="critical",
                description=f"Response only {len(response)} chars (threshold: {self.EXTREMELY_SHORT_THRESHOLD})",
                query_preview=query_preview,
                response_preview=resp_preview,
                metadata={"length": len(response)},
            ))
        elif len(response) < self.VERY_SHORT_THRESHOLD:
            issues.append(ResponseIssue(
                issue_type="very_short",
                severity="warning",
                description=f"Response only {len(response)} chars (threshold: {self.VERY_SHORT_THRESHOLD})",
                query_preview=query_preview,
                response_preview=resp_preview,
                metadata={"length": len(response)},
            ))

        # Truncation detection - ends mid-word
        if len(response) > 10 and response[-1].isalpha():
            # Check if it's not a valid single-letter answer
            last_word = response.split()[-1] if response.split() else ""
            if len(last_word) > 1:
                issues.append(ResponseIssue(
                    issue_type="truncated_mid_word",
                    severity="critical",
                    description=f"Response ends mid-word: '...{response[-20:]}'",
                    query_preview=query_preview,
                    response_preview=resp_preview,
                    metadata={"last_chars": response[-20:]},
                ))

        # Truncation detection - ends mid-sentence (not at valid ending)
        if len(response) > 20 and response[-1] not in self.VALID_ENDINGS:
            # Check it's not just ending with a letter (already caught above)
            if not response[-1].isalpha():
                issues.append(ResponseIssue(
                    issue_type="truncated_mid_sentence",
                    severity="warning",
                    description=f"Response may be truncated, ends with: '{response[-1]}'",
                    query_preview=query_preview,
                    response_preview=resp_preview,
                ))

        # Starts incomplete (lowercase, suggests truncation from start)
        if response[0].islower() and not response.startswith(("http", "www", "i ", "i'")):
            issues.append(ResponseIssue(
                issue_type="starts_incomplete",
                severity="info",
                description="Response starts with lowercase (may be truncated from start)",
                query_preview=query_preview,
                response_preview=resp_preview,
            ))

        # Contains error messages
        if self._error_regex.search(response):
            match = self._error_regex.search(response)
            issues.append(ResponseIssue(
                issue_type="contains_error",
                severity="warning",
                description=f"Response contains error pattern: '{match.group() if match else 'unknown'}'",
                query_preview=query_preview,
                response_preview=resp_preview,
            ))

        # Latency anomalies
        if latency < self.SUSPICIOUSLY_FAST_LATENCY and len(response) > 100:
            issues.append(ResponseIssue(
                issue_type="suspiciously_fast",
                severity="info",
                description=f"Response latency {latency:.2f}s seems too fast for {len(response)} chars",
                query_preview=query_preview,
                metadata={"latency": latency, "length": len(response)},
            ))

        if latency > self.VERY_SLOW_LATENCY:
            issues.append(ResponseIssue(
                issue_type="very_slow",
                severity="info",
                description=f"Response latency {latency:.2f}s exceeds {self.VERY_SLOW_LATENCY}s",
                query_preview=query_preview,
                metadata={"latency": latency},
            ))

        # Benchmark-specific checks
        issues.extend(self._check_benchmark_specific(response, task_id, query_preview))

        return issues

    def _check_benchmark_specific(
        self, response: str, task_id: str, query_preview: str
    ) -> list[ResponseIssue]:
        """Run benchmark-specific validation."""
        issues = []

        if not task_id:
            return issues

        task_lower = task_id.lower()

        # GSM8K: Should contain a numeric answer
        if "gsm8k" in task_lower:
            if not re.search(r"\d+", response):
                issues.append(ResponseIssue(
                    issue_type="gsm8k_no_number",
                    severity="warning",
                    description="GSM8K response contains no numeric answer",
                    query_preview=query_preview,
                    response_preview=response[:100],
                ))

        # Multiple choice: Should contain a valid option
        if any(x in task_lower for x in ["mmlu", "arc", "agieval", "truthfulqa"]):
            # Check for A/B/C/D/E patterns
            if not re.search(r"\b[A-E]\b|answer.{0,20}[A-E]", response, re.IGNORECASE):
                issues.append(ResponseIssue(
                    issue_type="mc_no_answer",
                    severity="info",
                    description="Multiple-choice response may not contain clear answer option",
                    query_preview=query_preview,
                    response_preview=response[:100],
                ))

        # BBEH: Various expected patterns
        if "bbeh" in task_lower:
            # Check for common BBEH answer formats
            valid_bbeh = any([
                re.search(r"(yes|no|true|false)\b", response, re.IGNORECASE),
                re.search(r"\b[A-E]\b", response),
                re.search(r"\d+", response),
            ])
            if not valid_bbeh and len(response) < 100:
                issues.append(ResponseIssue(
                    issue_type="bbeh_unclear_answer",
                    severity="info",
                    description="BBEH response may not contain clear answer",
                    query_preview=query_preview,
                    response_preview=response[:100],
                ))

        return issues


class DatasetQualityChecker:
    """Checks dataset-level quality issues."""

    def __init__(self, df: pd.DataFrame):
        """Initialize with router data DataFrame.

        Args:
            df: Router data with columns: query, llm, effect, task_id, etc.
        """
        self.df = df
        self.llms = df["llm"].unique().tolist()

    def check_duplicates(self) -> list[ResponseIssue]:
        """Check for duplicate query-LLM pairs."""
        issues = []

        dups = self.df.groupby(["query", "llm"]).size()
        dups = dups[dups > 1]

        if len(dups) > 0:
            # Check for inconsistent duplicates
            inconsistent = 0
            for (query, llm), count in dups.items():
                subset = self.df[(self.df["query"] == query) & (self.df["llm"] == llm)]
                if subset["effect"].nunique() > 1:
                    inconsistent += 1

            issues.append(ResponseIssue(
                issue_type="duplicate_pairs",
                severity="warning" if inconsistent == 0 else "critical",
                description=f"{len(dups)} duplicate query-LLM pairs ({inconsistent} with inconsistent effects)",
                metadata={"total_dups": len(dups), "inconsistent": inconsistent},
            ))

        return issues

    def check_effect_distribution(self) -> list[ResponseIssue]:
        """Check for anomalies in effect (performance) distribution."""
        issues = []

        for llm in self.llms:
            subset = self.df[self.df["llm"] == llm]["effect"]

            # Check for invalid values
            invalid = ((subset < 0) | (subset > 1)).sum()
            if invalid > 0:
                issues.append(ResponseIssue(
                    issue_type="invalid_effect",
                    severity="critical",
                    description=f"{llm}: {invalid} effect values outside [0, 1]",
                ))

            # Check for suspiciously uniform distribution
            if subset.std() < 0.1:
                issues.append(ResponseIssue(
                    issue_type="low_variance_effect",
                    severity="warning",
                    description=f"{llm}: Very low effect variance ({subset.std():.4f})",
                ))

            # Check for missing values
            missing = subset.isna().sum()
            if missing > 0:
                issues.append(ResponseIssue(
                    issue_type="missing_effect",
                    severity="critical",
                    description=f"{llm}: {missing} missing effect values",
                ))

        return issues

    def check_coverage(self) -> list[ResponseIssue]:
        """Check that all LLMs cover all queries."""
        issues = []

        queries_per_llm = self.df.groupby("llm")["query"].nunique()
        expected = queries_per_llm.max()

        for llm, count in queries_per_llm.items():
            if count < expected:
                issues.append(ResponseIssue(
                    issue_type="incomplete_coverage",
                    severity="warning",
                    description=f"{llm}: Only {count}/{expected} queries covered",
                    metadata={"expected": expected, "actual": count},
                ))

        return issues

    def check_cross_model_consistency(self) -> list[ResponseIssue]:
        """Check for suspicious cross-model patterns."""
        issues = []

        # Create pivot table
        pivot = self.df.pivot_table(
            index="query", columns="llm", values="effect", aggfunc="mean"
        )

        # Check correlations
        for i, llm1 in enumerate(self.llms):
            for llm2 in self.llms[i + 1:]:
                if llm1 in pivot.columns and llm2 in pivot.columns:
                    corr = pivot[llm1].corr(pivot[llm2])
                    if corr < 0.2:
                        issues.append(ResponseIssue(
                            issue_type="low_model_correlation",
                            severity="info",
                            description=f"Low correlation between {llm1} and {llm2}: {corr:.3f}",
                            metadata={"llm1": llm1, "llm2": llm2, "correlation": corr},
                        ))

        # Check for models with unexpectedly similar performance
        means = self.df.groupby("llm")["effect"].mean()

        # Gemini Flash vs Pro check
        if "Gemini 3 Flash Preview" in means.index and "Gemini 3 Pro Preview" in means.index:
            flash = means["Gemini 3 Flash Preview"]
            pro = means["Gemini 3 Pro Preview"]
            if flash > pro:
                issues.append(ResponseIssue(
                    issue_type="unexpected_ranking",
                    severity="warning",
                    description=f"Gemini 3 Flash ({flash:.4f}) outperforms Pro ({pro:.4f}) - unexpected",
                    metadata={"flash": flash, "pro": pro},
                ))

        return issues

    def check_task_distribution(self) -> list[ResponseIssue]:
        """Check for task distribution anomalies."""
        issues = []

        if "task_id" not in self.df.columns:
            return issues

        task_counts = self.df.groupby("task_id").size()
        expected_per_task = len(self.llms)

        for task, count in task_counts.items():
            if count % expected_per_task != 0:
                issues.append(ResponseIssue(
                    issue_type="task_count_mismatch",
                    severity="warning",
                    description=f"Task {task}: {count} rows not divisible by {expected_per_task} LLMs",
                ))

        return issues


class BenchmarkAlignmentChecker:
    """Validates that model scores align with public benchmark expectations.

    Checks:
    1. Best scores are within expected ranges for each task
    2. Tier 1 models outperform Tier 3 on reasoning tasks
    3. No suspicious ceiling effects (all models at 99%+)
    4. Cost-performance correlation is non-negative
    """

    def __init__(self, df: pd.DataFrame):
        """Initialize with router data DataFrame.

        Args:
            df: Router data with columns: task_id, llm, effect.
        """
        self.df = df
        self.pivot = df.pivot_table(
            index="task_id", columns="llm", values="effect", aggfunc="mean"
        )

    def check_benchmark_ranges(self) -> list[ResponseIssue]:
        """Check that best scores for each task are within expected ranges."""
        issues = []

        for task_id, (min_exp, max_exp) in PUBLIC_BENCHMARK_RANGES.items():
            if task_id not in self.pivot.index:
                continue

            best_score = self.pivot.loc[task_id].max()
            best_llm = self.pivot.loc[task_id].idxmax()

            if best_score < min_exp:
                issues.append(ResponseIssue(
                    issue_type="below_benchmark",
                    severity="warning",
                    description=(
                        f"{task_id}: best={best_score:.3f} ({best_llm}) < "
                        f"expected_min={min_exp}"
                    ),
                    metadata={
                        "task_id": task_id,
                        "best_score": best_score,
                        "best_llm": best_llm,
                        "expected_min": min_exp,
                    },
                ))
            elif best_score > max_exp:
                issues.append(ResponseIssue(
                    issue_type="ceiling_effect",
                    severity="warning",
                    description=(
                        f"{task_id}: best={best_score:.3f} ({best_llm}) > "
                        f"expected_max={max_exp} (data leak?)"
                    ),
                    metadata={
                        "task_id": task_id,
                        "best_score": best_score,
                        "best_llm": best_llm,
                        "expected_max": max_exp,
                    },
                ))

        return issues

    def check_capability_ordering(self) -> list[ResponseIssue]:
        """Check that frontier models outperform budget models on reasoning."""
        issues = []

        reasoning_tasks = [
            "GSM8K", "agieval", "mmlu_pro_chemistry", "mmlu_pro_law",
            "bbeh_multistep_arithmetic", "bbeh_boolean_expressions",
        ]

        tier1_llms = [k for k, v in LLM_CAPABILITY_TIERS.items() if v == 1]
        tier3_llms = [k for k, v in LLM_CAPABILITY_TIERS.items() if v == 3]

        tier1_cols = [c for c in tier1_llms if c in self.pivot.columns]
        tier3_cols = [c for c in tier3_llms if c in self.pivot.columns]

        if not tier1_cols or not tier3_cols:
            return issues

        violations = []
        for task in reasoning_tasks:
            if task not in self.pivot.index:
                continue

            tier1_mean = self.pivot.loc[task, tier1_cols].mean()
            tier3_mean = self.pivot.loc[task, tier3_cols].mean()

            if tier3_mean > tier1_mean:
                violations.append(
                    f"{task}: tier3={tier3_mean:.3f} > tier1={tier1_mean:.3f}"
                )

        if len(violations) > len(reasoning_tasks) // 2:
            issues.append(ResponseIssue(
                issue_type="tier_ordering_violation",
                severity="critical",
                description=(
                    f"Tier 3 (budget) beats Tier 1 (frontier) on "
                    f"{len(violations)}/{len(reasoning_tasks)} reasoning tasks"
                ),
                metadata={"violations": violations},
            ))

        return issues

    def check_suspicious_zeros(self) -> list[ResponseIssue]:
        """Check for LLMs with zero scores on enabled tasks (API failures)."""
        issues = []

        enabled_tasks = [t for t in self.pivot.index if t not in DISABLED_TASKS]
        enabled_matrix = self.pivot.loc[enabled_tasks]

        for llm in enabled_matrix.columns:
            zero_tasks = enabled_matrix.index[enabled_matrix[llm] == 0.0].tolist()
            if zero_tasks:
                issues.append(ResponseIssue(
                    issue_type="zero_scores",
                    severity="critical",
                    description=f"{llm}: {len(zero_tasks)} zero-score tasks",
                    metadata={"llm": llm, "zero_tasks": zero_tasks},
                ))

        return issues

    def check_cost_performance_correlation(self) -> list[ResponseIssue]:
        """Check that cost correlates positively with performance."""
        issues = []

        if "cost" not in self.df.columns:
            return issues

        llm_stats = self.df.groupby("llm").agg({
            "effect": "mean",
            "cost": "mean",
        })

        correlation = llm_stats["effect"].corr(llm_stats["cost"])

        if correlation < 0:
            issues.append(ResponseIssue(
                issue_type="negative_cost_correlation",
                severity="critical",
                description=(
                    f"Negative cost-performance correlation ({correlation:.3f}) - "
                    "cheap models outperforming expensive ones"
                ),
                metadata={"correlation": correlation},
            ))

        return issues

    def get_ranking_summary(self) -> dict[str, Any]:
        """Get a summary of model rankings across tasks."""
        llm_means = self.pivot.mean().sort_values(ascending=False)

        return {
            "overall_ranking": llm_means.to_dict(),
            "best_model": llm_means.index[0],
            "worst_model": llm_means.index[-1],
            "spread": llm_means.max() - llm_means.min(),
        }

    def validate_all(self) -> list[ResponseIssue]:
        """Run all benchmark alignment checks."""
        issues = []
        issues.extend(self.check_benchmark_ranges())
        issues.extend(self.check_capability_ordering())
        issues.extend(self.check_suspicious_zeros())
        issues.extend(self.check_cost_performance_correlation())
        return issues


class DataQualityValidator:
    """Main validator that orchestrates all quality checks."""

    def __init__(
        self,
        data_path: str | Path | None = None,
        cache_path: str | Path | None = None,
    ):
        """Initialize validator.

        Args:
            data_path: Path to router_data.csv. Defaults to standard location.
            cache_path: Path to response cache JSON. Defaults to standard location.
        """
        self.data_path = Path(data_path) if data_path else (
            _MODULE_DIR / "data" / "router_data.csv"
        )
        self.cache_path = Path(cache_path) if cache_path else (
            _MODULE_DIR / "data" / ".router_cache" / "response_cache.json"
        )

        self.response_checker = ResponseQualityChecker()
        self.df: pd.DataFrame | None = None
        self.cache: dict | None = None

    def load_data(self) -> None:
        """Load router data and cache."""
        logger.info("Loading router data from %s", self.data_path)
        self.df = pd.read_csv(self.data_path)
        self.df = self.df.drop_duplicates(subset=["query", "llm"])
        logger.info("Loaded %d rows", len(self.df))

        if self.cache_path.exists():
            logger.info("Loading response cache from %s", self.cache_path)
            with open(self.cache_path) as f:
                self.cache = json.load(f)
            logger.info("Loaded %d cached responses", len(self.cache))
        else:
            logger.warning("Cache not found at %s", self.cache_path)
            self.cache = {}

    def validate(
        self,
        llms: list[str] | None = None,
        quick: bool = False,
        sample_size: int | None = None,
    ) -> DataQualityReport:
        """Run full validation.

        Args:
            llms: Specific LLMs to validate. None = all.
            quick: Skip expensive cross-model checks.
            sample_size: Limit number of queries per LLM (for faster runs).

        Returns:
            Complete data quality report.
        """
        if self.df is None:
            self.load_data()

        report = DataQualityReport()

        # Dataset-level checks
        logger.info("Running dataset-level checks...")
        dataset_checker = DatasetQualityChecker(self.df)

        report.global_issues.extend(dataset_checker.check_duplicates())
        report.global_issues.extend(dataset_checker.check_effect_distribution())
        report.global_issues.extend(dataset_checker.check_coverage())
        report.global_issues.extend(dataset_checker.check_task_distribution())

        if not quick:
            report.global_issues.extend(dataset_checker.check_cross_model_consistency())

        # Benchmark alignment checks
        logger.info("Running benchmark alignment checks...")
        benchmark_checker = BenchmarkAlignmentChecker(self.df)
        report.global_issues.extend(benchmark_checker.validate_all())
        report.summary["ranking"] = benchmark_checker.get_ranking_summary()

        # Per-LLM response checks
        target_llms = llms or self.df["llm"].unique().tolist()

        for llm in target_llms:
            logger.info("Validating %s...", llm)
            llm_report = self._validate_llm(llm, sample_size)
            report.llm_reports[llm] = llm_report

        # Generate summary
        report.summary = self._generate_summary(report)

        return report

    def _validate_llm(
        self, llm: str, sample_size: int | None = None
    ) -> LLMQualityReport:
        """Validate a single LLM's responses."""
        report = LLMQualityReport(llm_name=llm)

        subset = self.df[self.df["llm"] == llm]
        if sample_size:
            subset = subset.head(sample_size)

        report.total_responses = len(subset)

        for _, row in subset.iterrows():
            query = row["query"]
            task_id = row.get("task_id", "")

            cache_key = _query_hash(query, llm)
            if cache_key in self.cache:
                cached = self.cache[cache_key]
                response = cached.get("response", "")
                latency = cached.get("latency", 0.0)

                issues = self.response_checker.check_response(
                    response=response,
                    latency=latency,
                    query=query,
                    task_id=task_id,
                )

                for issue in issues:
                    report.add_issue(issue)
            else:
                report.add_issue(ResponseIssue(
                    issue_type="missing_cache",
                    severity="warning",
                    description="Response not found in cache",
                    query_preview=query[:100],
                ))

        # Compute metrics
        report.metrics["effect_mean"] = subset["effect"].mean()
        report.metrics["effect_std"] = subset["effect"].std()
        report.metrics["accuracy"] = (subset["effect"] >= 0.5).mean()

        return report

    def _generate_summary(self, report: DataQualityReport) -> dict:
        """Generate summary statistics."""
        total_issues = sum(
            len(r.issues) for r in report.llm_reports.values()
        ) + len(report.global_issues)

        critical_issues = sum(
            sum(1 for i in r.issues if i.severity == "critical")
            for r in report.llm_reports.values()
        ) + sum(1 for i in report.global_issues if i.severity == "critical")

        # Identify worst LLMs by issue rate
        worst_by_truncation = max(
            report.llm_reports.items(),
            key=lambda x: x[1].get_issue_rate("truncated_mid_word"),
            default=(None, None),
        )

        worst_by_short = max(
            report.llm_reports.items(),
            key=lambda x: x[1].get_issue_rate("very_short") + x[1].get_issue_rate("extremely_short"),
            default=(None, None),
        )

        return {
            "total_issues": total_issues,
            "critical_issues": critical_issues,
            "llms_validated": len(report.llm_reports),
            "worst_truncation": {
                "llm": worst_by_truncation[0],
                "rate": worst_by_truncation[1].get_issue_rate("truncated_mid_word") if worst_by_truncation[1] else 0,
            } if worst_by_truncation[0] else None,
            "worst_short_responses": {
                "llm": worst_by_short[0],
                "rate": (
                    worst_by_short[1].get_issue_rate("very_short") +
                    worst_by_short[1].get_issue_rate("extremely_short")
                ) if worst_by_short[1] else 0,
            } if worst_by_short[0] else None,
            "recommendations": self._generate_recommendations(report),
        }

    def _generate_recommendations(self, report: DataQualityReport) -> list[str]:
        """Generate actionable recommendations."""
        recommendations = []

        # Check for high truncation rates
        for llm, llm_report in report.llm_reports.items():
            trunc_rate = llm_report.get_issue_rate("truncated_mid_word")
            if trunc_rate > 10:
                recommendations.append(
                    f"REGENERATE {llm}: {trunc_rate:.1f}% truncated responses"
                )

            short_rate = (
                llm_report.get_issue_rate("very_short") +
                llm_report.get_issue_rate("extremely_short")
            )
            if short_rate > 10:
                recommendations.append(
                    f"INVESTIGATE {llm}: {short_rate:.1f}% very short responses"
                )

            error_rate = llm_report.get_issue_rate("contains_error")
            if error_rate > 3:
                recommendations.append(
                    f"CHECK API for {llm}: {error_rate:.1f}% contain error messages"
                )

        # Global issues
        for issue in report.global_issues:
            if issue.issue_type == "unexpected_ranking":
                recommendations.append(
                    f"VALIDATE: {issue.description}"
                )
            if issue.issue_type == "duplicate_pairs" and issue.metadata.get("inconsistent", 0) > 0:
                recommendations.append(
                    f"FIX DUPLICATES: {issue.metadata['inconsistent']} pairs have inconsistent effects"
                )

        return recommendations

    def print_report(self, report: DataQualityReport) -> None:
        """Print human-readable report to console."""
        print("\n" + "=" * 80)
        print("DATA QUALITY VALIDATION REPORT")
        print("=" * 80)

        # Summary
        print(f"\nTotal issues found: {report.summary['total_issues']}")
        print(f"Critical issues: {report.summary['critical_issues']}")
        print(f"LLMs validated: {report.summary['llms_validated']}")

        # Model ranking (if available)
        if "ranking" in report.summary:
            ranking = report.summary["ranking"]
            print("\n" + "-" * 80)
            print("MODEL RANKING (by mean effect)")
            print("-" * 80)
            for i, (llm, score) in enumerate(
                ranking.get("overall_ranking", {}).items(), 1
            ):
                tier = LLM_CAPABILITY_TIERS.get(llm, "?")
                print(f"  {i}. {llm:<30} {score:.4f}  (tier {tier})")
            print(f"\n  Spread: {ranking.get('spread', 0):.4f}")

        # Per-LLM summary
        print("\n" + "-" * 80)
        print("PER-LLM ISSUE RATES")
        print("-" * 80)
        print(f"{'LLM':<30} {'Total':>8} {'Trunc%':>8} {'Short%':>8} {'Error%':>8}")
        print("-" * 80)

        for llm, llm_report in sorted(report.llm_reports.items()):
            trunc = llm_report.get_issue_rate("truncated_mid_word")
            short = llm_report.get_issue_rate("very_short") + llm_report.get_issue_rate("extremely_short")
            error = llm_report.get_issue_rate("contains_error")
            print(f"{llm:<30} {llm_report.total_responses:>8} {trunc:>7.1f}% {short:>7.1f}% {error:>7.1f}%")

        # Global issues
        if report.global_issues:
            print("\n" + "-" * 80)
            print("GLOBAL ISSUES")
            print("-" * 80)
            for issue in report.global_issues:
                severity_marker = "!!!" if issue.severity == "critical" else "!" if issue.severity == "warning" else ""
                print(f"  [{issue.severity.upper()}] {severity_marker} {issue.description}")

        # Recommendations
        if report.summary.get("recommendations"):
            print("\n" + "-" * 80)
            print("RECOMMENDATIONS")
            print("-" * 80)
            for rec in report.summary["recommendations"]:
                print(f"  * {rec}")

        print("\n" + "=" * 80)


class BadResponseExporter:
    """Export bad responses for regeneration."""

    # Issue types that warrant regeneration
    CRITICAL_ISSUES = {
        "empty_response",
        "extremely_short",
        "truncated_mid_word",
    }

    REGENERATE_ISSUES = CRITICAL_ISSUES | {
        "very_short",
        "contains_error",
    }

    def __init__(
        self,
        data_path: str | Path | None = None,
        cache_path: str | Path | None = None,
    ):
        self.data_path = Path(data_path) if data_path else (
            _MODULE_DIR / "data" / "router_data.csv"
        )
        self.cache_path = Path(cache_path) if cache_path else (
            _MODULE_DIR / "data" / ".router_cache" / "response_cache.json"
        )
        self.response_checker = ResponseQualityChecker()

    def find_bad_responses(
        self,
        llms: list[str] | None = None,
        issue_types: set[str] | None = None,
        min_severity: str = "warning",
    ) -> pd.DataFrame:
        """Find all responses with quality issues.

        Args:
            llms: Filter to specific LLMs. None = all.
            issue_types: Filter to specific issue types. None = all regenerate-worthy.
            min_severity: Minimum severity to include ("info", "warning", "critical").

        Returns:
            DataFrame with columns: query, llm, task_id, issue_type, severity, cache_key
        """
        logger.info("Loading data...")
        df = pd.read_csv(self.data_path)
        df = df.drop_duplicates(subset=["query", "llm"])

        with open(self.cache_path) as f:
            cache = json.load(f)

        target_issues = issue_types or self.REGENERATE_ISSUES
        severity_order = {"info": 0, "warning": 1, "critical": 2}
        min_sev_val = severity_order.get(min_severity, 1)

        if llms:
            df = df[df["llm"].isin(llms)]

        bad_responses = []
        logger.info("Scanning %d responses...", len(df))

        for _, row in df.iterrows():
            query = row["query"]
            llm = row["llm"]
            task_id = row.get("task_id", "")
            cache_key = _query_hash(query, llm)

            if cache_key not in cache:
                bad_responses.append({
                    "query": query,
                    "llm": llm,
                    "task_id": task_id,
                    "issue_type": "missing_cache",
                    "severity": "critical",
                    "cache_key": cache_key,
                })
                continue

            cached = cache[cache_key]
            response = cached.get("response", "")
            latency = cached.get("latency", 0.0)

            issues = self.response_checker.check_response(
                response=response,
                latency=latency,
                query=query,
                task_id=task_id,
            )

            for issue in issues:
                if issue.issue_type in target_issues:
                    if severity_order.get(issue.severity, 0) >= min_sev_val:
                        bad_responses.append({
                            "query": query,
                            "llm": llm,
                            "task_id": task_id,
                            "issue_type": issue.issue_type,
                            "severity": issue.severity,
                            "cache_key": cache_key,
                        })
                        break  # One issue per response is enough

        result = pd.DataFrame(bad_responses)
        logger.info("Found %d bad responses", len(result))
        return result

    def export_for_regeneration(
        self,
        output_path: str | Path,
        llms: list[str] | None = None,
        issue_types: set[str] | None = None,
    ) -> int:
        """Export bad responses to CSV for regeneration.

        Args:
            output_path: Where to save the CSV.
            llms: Filter to specific LLMs.
            issue_types: Filter to specific issue types.

        Returns:
            Number of responses exported.
        """
        bad_df = self.find_bad_responses(llms=llms, issue_types=issue_types)

        if bad_df.empty:
            logger.info("No bad responses found!")
            return 0

        output_path = Path(output_path)
        bad_df.to_csv(output_path, index=False)
        logger.info("Exported %d bad responses to %s", len(bad_df), output_path)

        # Print summary
        print("\n--- Bad Responses by LLM ---")
        print(bad_df.groupby("llm").size().sort_values(ascending=False))

        print("\n--- Bad Responses by Issue Type ---")
        print(bad_df.groupby("issue_type").size().sort_values(ascending=False))

        return len(bad_df)

    def remove_bad_from_cache(
        self,
        llms: list[str] | None = None,
        issue_types: set[str] | None = None,
        dry_run: bool = True,
    ) -> int:
        """Remove bad responses from cache to force regeneration.

        Args:
            llms: Filter to specific LLMs.
            issue_types: Filter to specific issue types.
            dry_run: If True, only report what would be removed.

        Returns:
            Number of responses removed (or would be removed).
        """
        bad_df = self.find_bad_responses(llms=llms, issue_types=issue_types)

        if bad_df.empty:
            logger.info("No bad responses to remove!")
            return 0

        cache_keys = bad_df["cache_key"].unique().tolist()
        logger.info("Found %d cache keys to remove", len(cache_keys))

        if dry_run:
            logger.info("DRY RUN - would remove %d cache entries", len(cache_keys))
            return len(cache_keys)

        # Load and modify cache
        with open(self.cache_path) as f:
            cache = json.load(f)

        removed = 0
        for key in cache_keys:
            if key in cache:
                del cache[key]
                removed += 1

        # Backup and save
        backup_path = self.cache_path.with_suffix(".json.backup")
        import shutil
        shutil.copy(self.cache_path, backup_path)
        logger.info("Backed up cache to %s", backup_path)

        with open(self.cache_path, "w") as f:
            json.dump(cache, f)
        logger.info("Removed %d bad responses from cache", removed)

        return removed


def main():
    """CLI entry point."""
    import argparse

    parser = argparse.ArgumentParser(
        description="Validate router training data quality"
    )
    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # Validate command
    validate_parser = subparsers.add_parser("validate", help="Run validation")
    validate_parser.add_argument(
        "--data", type=str, help="Path to router_data.csv"
    )
    validate_parser.add_argument(
        "--cache", type=str, help="Path to response cache JSON"
    )
    validate_parser.add_argument(
        "--llms", nargs="+", help="Specific LLMs to validate"
    )
    validate_parser.add_argument(
        "--quick", action="store_true", help="Skip expensive cross-model checks"
    )
    validate_parser.add_argument(
        "--sample", type=int, help="Sample size per LLM (for faster runs)"
    )
    validate_parser.add_argument(
        "--output", type=str, help="Export report to JSON file"
    )

    # Export bad responses command
    export_parser = subparsers.add_parser("export-bad", help="Export bad responses")
    export_parser.add_argument(
        "--output", type=str, required=True, help="Output CSV path"
    )
    export_parser.add_argument(
        "--llms", nargs="+", help="Specific LLMs to check"
    )
    export_parser.add_argument(
        "--critical-only", action="store_true", help="Only critical issues"
    )

    # Clean cache command
    clean_parser = subparsers.add_parser("clean-cache", help="Remove bad responses from cache")
    clean_parser.add_argument(
        "--llms", nargs="+", help="Specific LLMs to clean"
    )
    clean_parser.add_argument(
        "--dry-run", action="store_true", help="Only show what would be removed"
    )
    clean_parser.add_argument(
        "--critical-only", action="store_true", help="Only remove critical issues"
    )

    # Benchmark alignment command
    benchmark_parser = subparsers.add_parser(
        "benchmark", help="Check alignment with public benchmarks"
    )
    benchmark_parser.add_argument(
        "--data", type=str, help="Path to router_data.csv"
    )

    args = parser.parse_args()

    # Default to validate if no command specified
    if args.command is None or args.command == "validate":
        validator = DataQualityValidator(
            data_path=getattr(args, "data", None),
            cache_path=getattr(args, "cache", None),
        )

        report = validator.validate(
            llms=getattr(args, "llms", None),
            quick=getattr(args, "quick", False),
            sample_size=getattr(args, "sample", None),
        )

        validator.print_report(report)

        output = getattr(args, "output", None)
        if output:
            output_path = Path(output)
            output_path.write_text(
                json.dumps(report.to_dict(), indent=2),
                encoding="utf-8",
            )
            logger.info("Report exported to %s", output_path)

    elif args.command == "export-bad":
        exporter = BadResponseExporter()
        issue_types = BadResponseExporter.CRITICAL_ISSUES if args.critical_only else None
        exporter.export_for_regeneration(
            output_path=args.output,
            llms=args.llms,
            issue_types=issue_types,
        )

    elif args.command == "clean-cache":
        exporter = BadResponseExporter()
        issue_types = BadResponseExporter.CRITICAL_ISSUES if args.critical_only else None
        exporter.remove_bad_from_cache(
            llms=args.llms,
            issue_types=issue_types,
            dry_run=args.dry_run,
        )

    elif args.command == "benchmark":
        data_path = Path(args.data) if args.data else (
            _MODULE_DIR / "data" / "router_data.csv"
        )
        logger.info("Loading data from %s", data_path)
        df = pd.read_csv(data_path)

        checker = BenchmarkAlignmentChecker(df)
        issues = checker.validate_all()
        ranking = checker.get_ranking_summary()

        print("\n" + "=" * 80)
        print("BENCHMARK ALIGNMENT REPORT")
        print("=" * 80)

        # Print ranking
        print("\nMODEL RANKING (by mean effect)")
        print("-" * 60)
        for i, (llm, score) in enumerate(
            ranking.get("overall_ranking", {}).items(), 1
        ):
            tier = LLM_CAPABILITY_TIERS.get(llm, "?")
            print(f"  {i}. {llm:<30} {score:.4f}  (tier {tier})")
        print(f"\n  Best: {ranking['best_model']}")
        print(f"  Worst: {ranking['worst_model']}")
        print(f"  Spread: {ranking['spread']:.4f}")

        # Print issues
        if issues:
            print("\n" + "-" * 60)
            print("ALIGNMENT ISSUES")
            print("-" * 60)
            for issue in issues:
                severity_icon = (
                    "!!!" if issue.severity == "critical"
                    else "!" if issue.severity == "warning"
                    else ""
                )
                print(f"  [{issue.severity.upper()}] {severity_icon} {issue.description}")
        else:
            print("\n✓ All benchmark checks passed!")

        print("\n" + "=" * 80)


if __name__ == "__main__":
    main()
