"""Test Impact Analysis - run only tests affected by code changes.

Based on Launchable/Meta's approach: run 33% of tests, catch 99.9% of regressions.

Process:
1. Build coverage map: test -> files it executes
2. Get changed files from git diff
3. Select tests that execute changed files
4. Run only selected tests

Usage:
    selector = SmartTestSelector()
    await selector.build_coverage_map("pytest")
    result = selector.select_tests(changed_files=["src/auth.py"])

    print(f"Running {len(result.selected)} of {len(result.all_tests)} tests")
    print(f"Estimated time saved: {result.estimated_time_saved_percent}%")
"""

from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional, Set


@dataclass
class TestSelectionResult:
    """Result of smart test selection."""
    selected_tests: List[str]
    skipped_tests: List[str]
    all_tests: List[str]
    changed_files: List[str]
    selection_reason: Dict[str, str]  # test -> why selected/skipped
    estimated_time_saved_percent: float
    confidence: float  # How confident we are in the selection
    coverage_freshness_days: float  # How old is the coverage data

    @property
    def selection_ratio(self) -> float:
        """Ratio of tests selected to total tests."""
        if not self.all_tests:
            return 1.0
        return len(self.selected_tests) / len(self.all_tests)

    def to_pytest_args(self) -> str:
        """Convert selected tests to pytest arguments."""
        if not self.selected_tests:
            return ""
        return " ".join(self.selected_tests)


@dataclass
class CoverageData:
    """Coverage mapping data."""
    test_to_files: Dict[str, Set[str]]  # test -> files it covers
    file_to_tests: Dict[str, Set[str]]  # file -> tests that cover it
    test_durations: Dict[str, float]  # test -> average duration
    collected_at: datetime
    total_tests: int
    total_files: int

    def get_tests_for_files(self, files: Set[str]) -> Set[str]:
        """Get all tests that cover any of the given files."""
        tests = set()
        for file in files:
            # Exact match
            if file in self.file_to_tests:
                tests.update(self.file_to_tests[file])
            # Also check partial matches for renamed/moved files
            for covered_file, covering_tests in self.file_to_tests.items():
                if file.endswith(covered_file) or covered_file.endswith(file):
                    tests.update(covering_tests)
        return tests

    @property
    def freshness_days(self) -> float:
        """Days since coverage was collected."""
        age = datetime.now(timezone.utc) - self.collected_at
        return age.total_seconds() / (24 * 3600)


class CoverageMapper:
    """Build and maintain test-to-file coverage mapping."""

    def __init__(
        self,
        repo_path: Optional[Path] = None,
        cache_path: Optional[Path] = None,
    ):
        """Initialize coverage mapper.

        Args:
            repo_path: Path to repository.
            cache_path: Path to store coverage cache. Defaults to .hikaflow/coverage.json
        """
        self.repo_path = repo_path or Path.cwd()
        self.cache_path = cache_path or (self.repo_path / ".hikaflow" / "coverage.json")
        self._coverage: Optional[CoverageData] = None

    async def build_coverage_map(
        self,
        test_command: str = "pytest",
        force_rebuild: bool = False,
    ) -> CoverageData:
        """Build coverage mapping by running tests with coverage.

        Args:
            test_command: Base test command (e.g., "pytest").
            force_rebuild: Force rebuild even if cache exists.

        Returns:
            CoverageData with test-to-file mapping.
        """
        # Check cache first
        if not force_rebuild and self.cache_path.exists():
            cached = self._load_cache()
            if cached and cached.freshness_days < 7:  # Cache valid for 7 days
                self._coverage = cached
                return cached

        # Run tests with coverage
        coverage_file = self.repo_path / ".coverage"

        # Use pytest-cov to generate per-test coverage
        command = f"{test_command} --cov=. --cov-context=test --cov-report=json:.coverage.json -q"

        process = await asyncio.create_subprocess_shell(
            command,
            cwd=self.repo_path,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        await process.communicate()

        # Parse coverage data
        coverage_json = self.repo_path / ".coverage.json"
        if not coverage_json.exists():
            # Fallback: try to get test list without coverage
            return await self._build_fallback_coverage(test_command)

        coverage_data = self._parse_coverage_json(coverage_json)

        # Cache it
        self._save_cache(coverage_data)
        self._coverage = coverage_data

        return coverage_data

    def _parse_coverage_json(self, coverage_json: Path) -> CoverageData:
        """Parse pytest-cov JSON output into CoverageData."""
        with open(coverage_json) as f:
            data = json.load(f)

        test_to_files: Dict[str, Set[str]] = {}
        file_to_tests: Dict[str, Set[str]] = {}

        # Coverage.py with contexts has format:
        # {"files": {"path": {"contexts": {"test_name": [lines]}}}}
        for file_path, file_data in data.get("files", {}).items():
            contexts = file_data.get("contexts", {})
            for context_name, lines in contexts.items():
                if not context_name or context_name == "":
                    continue

                # Normalize test name (remove parameterization)
                test_name = context_name.split("[")[0]

                if test_name not in test_to_files:
                    test_to_files[test_name] = set()
                test_to_files[test_name].add(file_path)

                if file_path not in file_to_tests:
                    file_to_tests[file_path] = set()
                file_to_tests[file_path].add(test_name)

        return CoverageData(
            test_to_files=test_to_files,
            file_to_tests=file_to_tests,
            test_durations={},  # Could be populated from pytest timing
            collected_at=datetime.now(timezone.utc),
            total_tests=len(test_to_files),
            total_files=len(file_to_tests),
        )

    async def _build_fallback_coverage(self, test_command: str) -> CoverageData:
        """Build a basic coverage map when full coverage isn't available.

        Uses heuristics: tests in tests/test_foo.py likely test src/foo.py
        """
        # Collect all test files
        test_files = list(self.repo_path.glob("**/test_*.py"))
        test_files.extend(self.repo_path.glob("**/*_test.py"))

        test_to_files: Dict[str, Set[str]] = {}
        file_to_tests: Dict[str, Set[str]] = {}

        for test_file in test_files:
            # Heuristic: test_foo.py likely tests foo.py
            test_name = test_file.stem
            if test_name.startswith("test_"):
                module_name = test_name[5:]  # Remove "test_" prefix
            elif test_name.endswith("_test"):
                module_name = test_name[:-5]  # Remove "_test" suffix
            else:
                continue

            # Find matching source files
            source_patterns = [
                f"**/{module_name}.py",
                f"**/src/{module_name}.py",
                f"**/lib/{module_name}.py",
            ]

            matched_files = set()
            for pattern in source_patterns:
                for source_file in self.repo_path.glob(pattern):
                    if "test" not in str(source_file):
                        matched_files.add(str(source_file.relative_to(self.repo_path)))

            if matched_files:
                test_key = str(test_file.relative_to(self.repo_path))
                test_to_files[test_key] = matched_files
                for f in matched_files:
                    if f not in file_to_tests:
                        file_to_tests[f] = set()
                    file_to_tests[f].add(test_key)

        return CoverageData(
            test_to_files=test_to_files,
            file_to_tests=file_to_tests,
            test_durations={},
            collected_at=datetime.now(timezone.utc),
            total_tests=len(test_to_files),
            total_files=len(file_to_tests),
        )

    def _load_cache(self) -> Optional[CoverageData]:
        """Load coverage data from cache."""
        try:
            with open(self.cache_path) as f:
                data = json.load(f)
            return CoverageData(
                test_to_files={k: set(v) for k, v in data["test_to_files"].items()},
                file_to_tests={k: set(v) for k, v in data["file_to_tests"].items()},
                test_durations=data.get("test_durations", {}),
                collected_at=datetime.fromisoformat(data["collected_at"]),
                total_tests=data["total_tests"],
                total_files=data["total_files"],
            )
        except (OSError, json.JSONDecodeError, KeyError):
            return None

    def _save_cache(self, coverage: CoverageData) -> None:
        """Save coverage data to cache."""
        self.cache_path.parent.mkdir(parents=True, exist_ok=True)
        data = {
            "test_to_files": {k: list(v) for k, v in coverage.test_to_files.items()},
            "file_to_tests": {k: list(v) for k, v in coverage.file_to_tests.items()},
            "test_durations": coverage.test_durations,
            "collected_at": coverage.collected_at.isoformat(),
            "total_tests": coverage.total_tests,
            "total_files": coverage.total_files,
        }
        with open(self.cache_path, "w") as f:
            json.dump(data, f, indent=2)


class SmartTestSelector:
    """Select only tests affected by code changes."""

    def __init__(
        self,
        mapper: Optional[CoverageMapper] = None,
        repo_path: Optional[Path] = None,
    ):
        """Initialize smart test selector.

        Args:
            mapper: CoverageMapper instance. Creates one if None.
            repo_path: Path to repository.
        """
        self.repo_path = repo_path or Path.cwd()
        self.mapper = mapper or CoverageMapper(repo_path)

    async def select_tests(
        self,
        base_ref: str = "main",
        head_ref: str = "HEAD",
        changed_files: Optional[List[str]] = None,
        include_new_tests: bool = True,
    ) -> TestSelectionResult:
        """Select tests to run based on changed files.

        Args:
            base_ref: Base git ref to compare against.
            head_ref: Head git ref (what changed).
            changed_files: Explicit list of changed files. Uses git diff if None.
            include_new_tests: Always include tests in changed test files.

        Returns:
            TestSelectionResult with selected and skipped tests.
        """
        # Get coverage data
        coverage = self.mapper._coverage
        if not coverage:
            coverage = await self.mapper.build_coverage_map()

        # Get changed files
        if changed_files is None:
            changed_files = await self._get_changed_files(base_ref, head_ref)

        changed_set = set(changed_files)

        # Get all known tests
        all_tests = list(coverage.test_to_files.keys())

        # Find tests that cover changed files
        affected_tests = coverage.get_tests_for_files(changed_set)

        # Always include tests from changed test files
        if include_new_tests:
            for f in changed_files:
                if "test" in f and f.endswith(".py"):
                    # This is a changed test file - include all its tests
                    for test in all_tests:
                        if f in test or test.startswith(f.replace(".py", "")):
                            affected_tests.add(test)

        # Build selection reasons
        selection_reason: Dict[str, str] = {}
        for test in all_tests:
            if test in affected_tests:
                # Find which file(s) triggered this test
                test_files = coverage.test_to_files.get(test, set())
                triggering = test_files & changed_set
                if triggering:
                    selection_reason[test] = f"covers: {', '.join(list(triggering)[:3])}"
                else:
                    selection_reason[test] = "in changed test file"
            else:
                selection_reason[test] = "no changed files covered"

        selected = list(affected_tests)
        skipped = [t for t in all_tests if t not in affected_tests]

        # Calculate time savings (rough estimate)
        total_duration = sum(coverage.test_durations.values()) or len(all_tests) * 1.0
        selected_duration = sum(coverage.test_durations.get(t, 1.0) for t in selected)
        time_saved = ((total_duration - selected_duration) / total_duration * 100) if total_duration else 0

        # Calculate confidence based on coverage freshness
        confidence = max(0.5, 1.0 - (coverage.freshness_days / 30))

        return TestSelectionResult(
            selected_tests=selected,
            skipped_tests=skipped,
            all_tests=all_tests,
            changed_files=changed_files,
            selection_reason=selection_reason,
            estimated_time_saved_percent=time_saved,
            confidence=confidence,
            coverage_freshness_days=coverage.freshness_days,
        )

    async def _get_changed_files(self, base_ref: str, head_ref: str) -> List[str]:
        """Get list of changed files between two refs."""
        process = await asyncio.create_subprocess_exec(
            "git", "diff", "--name-only", f"{base_ref}...{head_ref}",
            cwd=self.repo_path,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        stdout, _ = await process.communicate()
        return [f for f in stdout.decode().strip().split("\n") if f]
