"""
Regression test suite for comparing model/prompt changes.

Provides:
- Test dataset management
- Baseline comparisons
- Quality regression detection
"""

import json
import hashlib
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

from .metrics import CheatSheetMetrics, calculate_bleu, calculate_rouge
from .rag_eval import RAGEvaluator, RetrievalMetrics
from ..infra.logging import get_logger
from ..infra.tracing import trace_function

logger = get_logger(__name__)


@dataclass
class TestCase:
    """A single regression test case."""
    id: str
    name: str
    description: str = ""
    
    # Input
    document_path: str = ""
    document_hash: str = ""  # For integrity checking
    
    # Expected outputs (gold standard)
    expected_concepts: List[str] = field(default_factory=list)
    expected_formulas: List[str] = field(default_factory=list)
    expected_topics: List[str] = field(default_factory=list)
    
    # Quality thresholds
    min_concept_count: int = 5
    min_formula_count: int = 3
    min_coverage: float = 0.6
    max_hallucination_rate: float = 0.1
    
    # Metadata
    tags: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "document_path": self.document_path,
            "document_hash": self.document_hash,
            "expected_concepts": self.expected_concepts,
            "expected_formulas": self.expected_formulas,
            "expected_topics": self.expected_topics,
            "min_concept_count": self.min_concept_count,
            "min_formula_count": self.min_formula_count,
            "min_coverage": self.min_coverage,
            "max_hallucination_rate": self.max_hallucination_rate,
            "tags": self.tags,
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "TestCase":
        return cls(**data)


@dataclass
class TestResult:
    """Result from a single test case execution."""
    test_case_id: str
    passed: bool
    
    # Actual outputs
    actual_concept_count: int = 0
    actual_formula_count: int = 0
    actual_algorithm_count: int = 0
    
    # Quality scores
    coverage_score: float = 0.0
    concept_overlap: float = 0.0
    formula_overlap: float = 0.0
    hallucination_rate: float = 0.0
    
    # Text quality
    bleu_score: float = 0.0
    rouge_scores: Dict[str, float] = field(default_factory=dict)
    
    # Failure reasons
    failures: List[str] = field(default_factory=list)
    
    # Timing
    duration_ms: float = 0.0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "test_case_id": self.test_case_id,
            "passed": self.passed,
            "actual_concept_count": self.actual_concept_count,
            "actual_formula_count": self.actual_formula_count,
            "actual_algorithm_count": self.actual_algorithm_count,
            "coverage_score": self.coverage_score,
            "concept_overlap": self.concept_overlap,
            "formula_overlap": self.formula_overlap,
            "hallucination_rate": self.hallucination_rate,
            "bleu_score": self.bleu_score,
            "rouge_scores": self.rouge_scores,
            "failures": self.failures,
            "duration_ms": self.duration_ms,
        }


@dataclass
class RegressionReport:
    """Complete regression test report."""
    run_id: str
    timestamp: str
    
    # Configuration
    model_name: str = ""
    prompt_version: str = ""
    
    # Results
    total_tests: int = 0
    passed_tests: int = 0
    failed_tests: int = 0
    
    # Individual results
    results: List[TestResult] = field(default_factory=list)
    
    # Aggregate metrics
    avg_coverage: float = 0.0
    avg_concept_overlap: float = 0.0
    avg_formula_overlap: float = 0.0
    
    # Comparison to baseline
    baseline_comparison: Dict[str, float] = field(default_factory=dict)
    is_regression: bool = False
    regression_details: List[str] = field(default_factory=list)
    
    @property
    def pass_rate(self) -> float:
        return self.passed_tests / self.total_tests if self.total_tests > 0 else 0.0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "run_id": self.run_id,
            "timestamp": self.timestamp,
            "model_name": self.model_name,
            "prompt_version": self.prompt_version,
            "total_tests": self.total_tests,
            "passed_tests": self.passed_tests,
            "failed_tests": self.failed_tests,
            "pass_rate": self.pass_rate,
            "results": [r.to_dict() for r in self.results],
            "avg_coverage": self.avg_coverage,
            "avg_concept_overlap": self.avg_concept_overlap,
            "avg_formula_overlap": self.avg_formula_overlap,
            "baseline_comparison": self.baseline_comparison,
            "is_regression": self.is_regression,
            "regression_details": self.regression_details,
        }


class RegressionTestSuite:
    """
    Regression test suite for quality assurance.
    
    Manages test cases, runs comparisons, and detects regressions.
    """
    
    def __init__(
        self,
        test_cases_path: Optional[Union[str, Path]] = None,
        baseline_path: Optional[Union[str, Path]] = None,
        output_dir: Union[str, Path] = "./regression_tests"
    ):
        """
        Initialize regression test suite.
        
        Args:
            test_cases_path: Path to test cases JSON file
            baseline_path: Path to baseline results JSON file
            output_dir: Directory for storing results
        """
        self._output_dir = Path(output_dir)
        self._output_dir.mkdir(parents=True, exist_ok=True)
        
        self._test_cases: List[TestCase] = []
        self._baseline: Optional[RegressionReport] = None
        
        if test_cases_path and Path(test_cases_path).exists():
            self._load_test_cases(Path(test_cases_path))
        
        if baseline_path and Path(baseline_path).exists():
            self._load_baseline(Path(baseline_path))
    
    def _load_test_cases(self, path: Path) -> None:
        """Load test cases from JSON file."""
        with open(path) as f:
            data = json.load(f)
        
        self._test_cases = [TestCase.from_dict(tc) for tc in data.get("test_cases", [])]
        logger.info(f"Loaded {len(self._test_cases)} test cases")
    
    def _load_baseline(self, path: Path) -> None:
        """Load baseline results from JSON file."""
        with open(path) as f:
            data = json.load(f)
        
        self._baseline = RegressionReport(
            run_id=data.get("run_id", ""),
            timestamp=data.get("timestamp", ""),
            model_name=data.get("model_name", ""),
            prompt_version=data.get("prompt_version", ""),
            total_tests=data.get("total_tests", 0),
            passed_tests=data.get("passed_tests", 0),
            failed_tests=data.get("failed_tests", 0),
            avg_coverage=data.get("avg_coverage", 0.0),
            avg_concept_overlap=data.get("avg_concept_overlap", 0.0),
            avg_formula_overlap=data.get("avg_formula_overlap", 0.0),
        )
        logger.info(f"Loaded baseline: {self._baseline.run_id}")
    
    def add_test_case(self, test_case: TestCase) -> None:
        """Add a test case to the suite."""
        self._test_cases.append(test_case)
    
    def save_test_cases(self, path: Union[str, Path]) -> None:
        """Save test cases to JSON file."""
        path = Path(path)
        data = {
            "test_cases": [tc.to_dict() for tc in self._test_cases],
            "last_updated": datetime.now().isoformat()
        }
        with open(path, "w") as f:
            json.dump(data, f, indent=2)
    
    @trace_function(name="run_regression_suite")
    def run(
        self,
        pipeline,
        model_name: str = "",
        prompt_version: str = "",
        tags: Optional[List[str]] = None
    ) -> RegressionReport:
        """
        Run the regression test suite.
        
        Args:
            pipeline: CheatSheetPipeline instance to test
            model_name: Model being tested
            prompt_version: Prompt version being tested
            tags: Optional tags to filter test cases
            
        Returns:
            RegressionReport with results
        """
        import time
        import uuid
        
        run_id = f"regression_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"
        
        # Filter test cases by tags
        test_cases = self._test_cases
        if tags:
            test_cases = [tc for tc in test_cases if any(t in tc.tags for t in tags)]
        
        logger.info(f"Running {len(test_cases)} regression tests")
        
        results = []
        for tc in test_cases:
            try:
                result = self._run_single_test(tc, pipeline)
                results.append(result)
            except Exception as e:
                logger.error(f"Test {tc.id} failed with error: {e}")
                results.append(TestResult(
                    test_case_id=tc.id,
                    passed=False,
                    failures=[f"Exception: {str(e)}"]
                ))
        
        # Build report
        report = self._build_report(
            run_id=run_id,
            model_name=model_name,
            prompt_version=prompt_version,
            results=results
        )
        
        # Compare to baseline
        if self._baseline:
            self._compare_to_baseline(report)
        
        # Save report
        self._save_report(report)
        
        return report
    
    def _run_single_test(
        self,
        test_case: TestCase,
        pipeline
    ) -> TestResult:
        """Run a single test case."""
        import time
        
        start_time = time.perf_counter()
        failures = []
        
        # Skip if document not available
        if test_case.document_path and not Path(test_case.document_path).exists():
            return TestResult(
                test_case_id=test_case.id,
                passed=False,
                failures=[f"Document not found: {test_case.document_path}"]
            )
        
        # Run pipeline (mock for testing)
        try:
            if hasattr(pipeline, 'build') and test_case.document_path:
                result = pipeline.build(test_case.document_path, compile_pdf=False)
                content = result.content if hasattr(result, 'content') else None
            else:
                content = None
        except Exception as e:
            return TestResult(
                test_case_id=test_case.id,
                passed=False,
                failures=[f"Pipeline error: {str(e)}"]
            )
        
        # Extract counts
        actual_concepts = len(content.concepts) if content else 0
        actual_formulas = len(content.formulas) if content else 0
        actual_algorithms = len(content.algorithms) if content else 0
        
        # Check thresholds
        if actual_concepts < test_case.min_concept_count:
            failures.append(f"Concept count {actual_concepts} < min {test_case.min_concept_count}")
        
        if actual_formulas < test_case.min_formula_count:
            failures.append(f"Formula count {actual_formulas} < min {test_case.min_formula_count}")
        
        # Calculate overlaps
        concept_overlap = self._calculate_overlap(
            [c.get("name", "") for c in (content.concepts if content else [])],
            test_case.expected_concepts
        )
        
        formula_overlap = self._calculate_overlap(
            [f.get("name", "") for f in (content.formulas if content else [])],
            test_case.expected_formulas
        )
        
        coverage = (concept_overlap + formula_overlap) / 2 if test_case.expected_concepts or test_case.expected_formulas else 1.0
        
        if coverage < test_case.min_coverage:
            failures.append(f"Coverage {coverage:.2f} < min {test_case.min_coverage}")
        
        duration_ms = (time.perf_counter() - start_time) * 1000
        
        return TestResult(
            test_case_id=test_case.id,
            passed=len(failures) == 0,
            actual_concept_count=actual_concepts,
            actual_formula_count=actual_formulas,
            actual_algorithm_count=actual_algorithms,
            coverage_score=coverage,
            concept_overlap=concept_overlap,
            formula_overlap=formula_overlap,
            failures=failures,
            duration_ms=duration_ms
        )
    
    def _calculate_overlap(
        self,
        actual: List[str],
        expected: List[str]
    ) -> float:
        """Calculate overlap between actual and expected items."""
        if not expected:
            return 1.0
        
        actual_lower = {a.lower().strip() for a in actual if a}
        expected_lower = {e.lower().strip() for e in expected if e}
        
        matches = len(actual_lower & expected_lower)
        return matches / len(expected_lower)
    
    def _build_report(
        self,
        run_id: str,
        model_name: str,
        prompt_version: str,
        results: List[TestResult]
    ) -> RegressionReport:
        """Build regression report from results."""
        passed = sum(1 for r in results if r.passed)
        failed = len(results) - passed
        
        avg_coverage = sum(r.coverage_score for r in results) / len(results) if results else 0.0
        avg_concept = sum(r.concept_overlap for r in results) / len(results) if results else 0.0
        avg_formula = sum(r.formula_overlap for r in results) / len(results) if results else 0.0
        
        return RegressionReport(
            run_id=run_id,
            timestamp=datetime.now().isoformat(),
            model_name=model_name,
            prompt_version=prompt_version,
            total_tests=len(results),
            passed_tests=passed,
            failed_tests=failed,
            results=results,
            avg_coverage=avg_coverage,
            avg_concept_overlap=avg_concept,
            avg_formula_overlap=avg_formula,
        )
    
    def _compare_to_baseline(self, report: RegressionReport) -> None:
        """Compare report to baseline and detect regressions."""
        baseline = self._baseline
        
        report.baseline_comparison = {
            "coverage_delta": report.avg_coverage - baseline.avg_coverage,
            "concept_overlap_delta": report.avg_concept_overlap - baseline.avg_concept_overlap,
            "formula_overlap_delta": report.avg_formula_overlap - baseline.avg_formula_overlap,
            "pass_rate_delta": report.pass_rate - baseline.pass_rate,
        }
        
        # Check for regressions (significant drops)
        regression_threshold = -0.05  # 5% drop
        
        if report.baseline_comparison["coverage_delta"] < regression_threshold:
            report.is_regression = True
            report.regression_details.append(
                f"Coverage dropped: {baseline.avg_coverage:.2f} -> {report.avg_coverage:.2f}"
            )
        
        if report.baseline_comparison["pass_rate_delta"] < regression_threshold:
            report.is_regression = True
            report.regression_details.append(
                f"Pass rate dropped: {baseline.pass_rate:.2f} -> {report.pass_rate:.2f}"
            )
    
    def _save_report(self, report: RegressionReport) -> None:
        """Save regression report to file."""
        report_path = self._output_dir / f"{report.run_id}.json"
        with open(report_path, "w") as f:
            json.dump(report.to_dict(), f, indent=2)
        logger.info(f"Saved regression report: {report_path}")
    
    def set_baseline(self, report: RegressionReport) -> None:
        """Set a report as the new baseline."""
        self._baseline = report
        baseline_path = self._output_dir / "baseline.json"
        with open(baseline_path, "w") as f:
            json.dump(report.to_dict(), f, indent=2)
        logger.info(f"Set new baseline: {report.run_id}")
