"""Flaky test quarantine management.

This module manages a quarantine file (.flaky-quarantine) that tracks
flaky tests that should be skipped in CI to prevent false failures.
"""

from __future__ import annotations

import json
import logging
import subprocess
import sys
import tempfile
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional, Set

QUARANTINE_FILE = ".flaky-quarantine"


@dataclass
class QuarantinedTest:
    """A test that has been quarantined due to flakiness."""

    nodeid: str  # pytest nodeid (e.g., tests/test_foo.py::test_bar)
    flakiness_score: float  # 0-100, higher = more flaky
    pass_rate: float  # 0-1, percentage of passes
    quarantined_at: str  # ISO timestamp
    reason: str  # Why it was quarantined
    failure_signatures: List[str] = field(default_factory=list)  # Unique failure types
    total_runs: int = 0
    pass_count: int = 0
    fail_count: int = 0
    last_checked: Optional[str] = None
    auto_quarantined: bool = True  # vs manually added


@dataclass
class QuarantineFile:
    """The .flaky-quarantine file format."""

    version: int = 1
    updated_at: str = ""
    tests: Dict[str, QuarantinedTest] = field(default_factory=dict)

    def add_test(self, test: QuarantinedTest) -> None:
        """Add or update a quarantined test."""
        self.tests[test.nodeid] = test
        self.updated_at = datetime.now(timezone.utc).isoformat()

    def remove_test(self, nodeid: str) -> bool:
        """Remove a test from quarantine. Returns True if removed."""
        if nodeid in self.tests:
            del self.tests[nodeid]
            self.updated_at = datetime.now(timezone.utc).isoformat()
            return True
        return False

    def is_quarantined(self, nodeid: str) -> bool:
        """Check if a test is quarantined."""
        return nodeid in self.tests

    def get_quarantined_nodeids(self) -> Set[str]:
        """Get all quarantined test nodeids."""
        return set(self.tests.keys())


def load_quarantine_file(path: Optional[Path] = None) -> QuarantineFile:
    """Load quarantine file from disk.

    Args:
        path: Path to quarantine file. If None, searches current directory.

    Returns:
        QuarantineFile instance (empty if file doesn't exist)
    """
    if path is None:
        path = Path.cwd() / QUARANTINE_FILE

    if not path.exists():
        return QuarantineFile()

    try:
        with open(path, encoding="utf-8") as f:
            data = json.load(f)

        qf = QuarantineFile(
            version=data.get("version", 1),
            updated_at=data.get("updated_at", ""),
        )

        for nodeid, test_data in data.get("tests", {}).items():
            qf.tests[nodeid] = QuarantinedTest(
                nodeid=nodeid,
                flakiness_score=test_data.get("flakiness_score", 0),
                pass_rate=test_data.get("pass_rate", 0),
                quarantined_at=test_data.get("quarantined_at", ""),
                reason=test_data.get("reason", ""),
                failure_signatures=test_data.get("failure_signatures", []),
                total_runs=test_data.get("total_runs", 0),
                pass_count=test_data.get("pass_count", 0),
                fail_count=test_data.get("fail_count", 0),
                last_checked=test_data.get("last_checked"),
                auto_quarantined=test_data.get("auto_quarantined", True),
            )

        return qf
    except (json.JSONDecodeError, KeyError) as e:
        print(f"Warning: Could not parse {path}: {e}", file=sys.stderr)
        return QuarantineFile()


def save_quarantine_file(qf: QuarantineFile, path: Optional[Path] = None) -> None:
    """Save quarantine file to disk.

    Args:
        qf: QuarantineFile to save
        path: Path to save to. If None, saves to current directory.
    """
    if path is None:
        path = Path.cwd() / QUARANTINE_FILE

    data = {
        "version": qf.version,
        "updated_at": qf.updated_at or datetime.now(timezone.utc).isoformat(),
        "tests": {
            nodeid: asdict(test)
            for nodeid, test in qf.tests.items()
        },
    }

    # Atomic write: write to temp file then rename to prevent corruption
    dir_path = path.parent if isinstance(path, Path) else Path(path).parent
    fd, tmp_path = tempfile.mkstemp(dir=str(dir_path), suffix=".tmp")
    try:
        with open(fd, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)
        Path(tmp_path).replace(path)
    except BaseException:
        Path(tmp_path).unlink(missing_ok=True)
        raise


@dataclass
class TestRunResult:
    """Result of running a single test."""

    nodeid: str
    passed: bool
    duration: float
    error_type: Optional[str] = None
    error_message: Optional[str] = None


@dataclass
class FlakyScanResult:
    """Result of scanning for flaky tests."""

    total_tests: int
    flaky_tests: List[QuarantinedTest]
    stable_tests: int
    scan_duration: float
    runs_per_test: int


def run_pytest_collect(test_path: str) -> List[str]:
    """Collect test nodeids using pytest.

    Args:
        test_path: Path to tests directory or file

    Returns:
        List of test nodeids
    """
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pytest", "--collect-only", "-q", "-q", "--", test_path],
            capture_output=True,
            text=True,
            timeout=120,
        )
    except subprocess.TimeoutExpired:
        return []

    if result.returncode != 0:
        return []

    nodeids = []
    for line in result.stdout.strip().split("\n"):
        stripped = line.strip()
        if "::" in stripped:
            # This is a test nodeid
            nodeids.append(stripped)

    return nodeids


def run_single_test(nodeid: str, timeout: int = 60) -> TestRunResult:
    """Run a single test and return the result.

    Args:
        nodeid: pytest nodeid
        timeout: Timeout in seconds

    Returns:
        TestRunResult
    """
    start = time.time()

    try:
        result = subprocess.run(
            [sys.executable, "-m", "pytest", "-x", "-v", "--tb=short", "--", nodeid],
            capture_output=True,
            text=True,
            timeout=timeout,
        )

        duration = time.time() - start
        passed = result.returncode == 0

        error_type = None
        error_message = None

        if not passed:
            # Try to extract error info from output
            for line in result.stdout.split("\n") + result.stderr.split("\n"):
                if "Error" in line or "Exception" in line:
                    parts = line.split(":")
                    if len(parts) >= 2:
                        error_type = parts[0].strip().split()[-1]
                        error_message = ":".join(parts[1:]).strip()[:200]
                        break

        return TestRunResult(
            nodeid=nodeid,
            passed=passed,
            duration=duration,
            error_type=error_type,
            error_message=error_message,
        )

    except subprocess.TimeoutExpired:
        return TestRunResult(
            nodeid=nodeid,
            passed=False,
            duration=timeout,
            error_type="TimeoutError",
            error_message=f"Test timed out after {timeout}s",
        )
    except Exception as e:
        return TestRunResult(
            nodeid=nodeid,
            passed=False,
            duration=time.time() - start,
            error_type=type(e).__name__,
            error_message=str(e),
        )


def _run_batch_and_parse(
    test_path: str,
    timeout: int,
    xdist_workers: int,
) -> Dict[str, TestRunResult]:
    """Run all tests in one pytest invocation and parse per-test results from JUnit XML.

    Instead of spawning one subprocess per test (25K+ processes for large suites),
    this runs pytest once with --junitxml and optionally pytest-xdist for parallelism.
    Reduces a 10+ hour scan to ~1-3 minutes per run.

    Args:
        test_path: Path to test file or directory
        timeout: Overall timeout for the entire run in seconds
        xdist_workers: Number of xdist workers (0 = no xdist)

    Returns:
        Dict mapping nodeid -> TestRunResult
    """
    import tempfile
    try:
        import defusedxml.ElementTree as ET
    except ImportError:
        import xml.etree.ElementTree as ET

    results: Dict[str, TestRunResult] = {}

    with tempfile.NamedTemporaryFile(suffix=".xml", delete=False) as f:
        xml_path = f.name

    try:
        cmd = [
            sys.executable, "-m", "pytest",
            f"--junitxml={xml_path}",
            "--tb=line",
            "-q",
            "--timeout=30",  # per-test timeout
            "--", test_path,
        ]

        if xdist_workers > 0:
            cmd.extend(["-n", str(xdist_workers), "--dist=loadfile"])

        proc = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=timeout,
        )

        # Parse JUnit XML for per-test results
        try:
            tree = ET.parse(xml_path)
            root = tree.getroot()

            for testcase in root.iter("testcase"):
                classname = testcase.get("classname", "")
                name = testcase.get("name", "")
                time_taken = float(testcase.get("time", "0"))

                # Reconstruct nodeid from JUnit classname + name
                # classname="tests.test_auth.TestPKCEGeneration" name="test_foo"
                # → nodeid="tests/test_auth.py::TestPKCEGeneration::test_foo"
                parts = classname.split(".")
                # Find where the file module ends and class begins
                # Module parts are lowercase (tests, test_auth), class parts are CamelCase
                file_parts = []
                class_parts = []
                for p in parts:
                    if not class_parts and (p.startswith("test_") or p == "tests" or p == "conftest" or p.islower()):
                        file_parts.append(p)
                    else:
                        class_parts.append(p)
                file_path = "/".join(file_parts) + ".py"
                if class_parts:
                    nodeid = f"{file_path}::{'::'.join(class_parts)}::{name}"
                else:
                    nodeid = f"{file_path}::{name}"

                # Check for failures/errors
                failure = testcase.find("failure")
                error = testcase.find("error")
                skipped = testcase.find("skipped")

                if skipped is not None:
                    continue  # Skip skipped tests

                passed = failure is None and error is None
                error_type = None
                error_message = None

                if failure is not None:
                    error_type = failure.get("type", "AssertionError")
                    error_message = (failure.get("message") or "")[:200]
                elif error is not None:
                    error_type = error.get("type", "Error")
                    error_message = (error.get("message") or "")[:200]

                results[nodeid] = TestRunResult(
                    nodeid=nodeid,
                    passed=passed,
                    duration=time_taken,
                    error_type=error_type,
                    error_message=error_message,
                )
        except ET.ParseError as exc:
            logging.getLogger(__name__).warning("JUnit XML parse error for %s: %s", test_path, exc)

    except subprocess.TimeoutExpired:
        logging.getLogger(__name__).warning("Test run timed out after %ds: %s", timeout, test_path)
    except (OSError, subprocess.SubprocessError) as exc:
        logging.getLogger(__name__).warning("Test execution failed for %s: %s", test_path, exc)
    finally:
        try:
            Path(xml_path).unlink(missing_ok=True)
        except OSError:
            pass

    return results


def scan_for_flaky_tests(
    test_path: str,
    runs: int = 5,
    parallel: int = 4,
    timeout: int = 60,
    flaky_threshold: float = 0.8,
    progress_callback: Optional[callable] = None,
) -> FlakyScanResult:
    """Scan tests for flakiness by running the full suite N times.

    Uses batch execution (one pytest invocation per run) with pytest-xdist
    for parallelism. This is ~100-500x faster than running each test as a
    separate subprocess.

    Old approach: 8571 tests x 3 runs = 25,713 subprocess spawns (~10+ hours)
    New approach: 3 full pytest runs with xdist (~3-5 minutes)

    Args:
        test_path: Path to tests directory or file
        runs: Number of times to run the suite (2-10)
        parallel: Number of xdist workers for parallelism within each run
        timeout: Timeout per individual run in seconds
        flaky_threshold: Pass rate below which a test is considered flaky (0-1)
        progress_callback: Optional callback(completed, total, current_test)

    Returns:
        FlakyScanResult with flaky tests identified
    """
    import os

    start_time = time.time()

    # Determine xdist worker count — use CPU cores if available
    cpu_count = os.cpu_count() or 4
    xdist_workers = min(parallel * 2, cpu_count)

    # Per-run timeout: scale with test suite size
    # Single file: 60s, directory: up to 300s
    is_file = os.path.isfile(test_path)
    run_timeout = timeout if is_file else max(timeout, 300)

    # Track results per test across all runs
    test_results: Dict[str, List[TestRunResult]] = defaultdict(list)
    total_tests_seen = 0

    for run_num in range(runs):
        if progress_callback:
            progress_callback(run_num, runs, f"run {run_num + 1}/{runs}")

        batch_results = _run_batch_and_parse(
            test_path,
            timeout=run_timeout,
            xdist_workers=xdist_workers,
        )

        for nodeid, result in batch_results.items():
            test_results[nodeid].append(result)

        if batch_results:
            total_tests_seen = max(total_tests_seen, len(batch_results))

    if not test_results:
        return FlakyScanResult(
            total_tests=0,
            flaky_tests=[],
            stable_tests=0,
            scan_duration=time.time() - start_time,
            runs_per_test=runs,
        )

    # Analyze results to find flaky tests
    flaky_tests = []
    stable_count = 0

    for nodeid, results in test_results.items():
        pass_count = sum(1 for r in results if r.passed)
        fail_count = len(results) - pass_count
        pass_rate = pass_count / len(results) if results else 0

        # Collect unique failure signatures
        failure_sigs = set()
        for r in results:
            if not r.passed and r.error_type:
                failure_sigs.add(r.error_type)

        # A test is flaky if it has mixed results and pass rate is below threshold
        is_flaky = 0 < pass_rate < flaky_threshold

        if is_flaky:
            # Calculate flakiness score (0-100)
            # Higher score = more unpredictable
            flakiness_score = (1 - abs(pass_rate - 0.5) * 2) * 100

            flaky_test = QuarantinedTest(
                nodeid=nodeid,
                flakiness_score=flakiness_score,
                pass_rate=pass_rate,
                quarantined_at=datetime.now(timezone.utc).isoformat(),
                reason=f"Flaky: {pass_count}/{len(results)} passes ({pass_rate*100:.0f}%)",
                failure_signatures=list(failure_sigs),
                total_runs=len(results),
                pass_count=pass_count,
                fail_count=fail_count,
                last_checked=datetime.now(timezone.utc).isoformat(),
                auto_quarantined=True,
            )
            flaky_tests.append(flaky_test)
        else:
            stable_count += 1

    if progress_callback:
        progress_callback(runs, runs, "analyzing results")

    # Sort by flakiness score (most flaky first)
    flaky_tests.sort(key=lambda t: t.flakiness_score, reverse=True)

    return FlakyScanResult(
        total_tests=total_tests_seen,
        flaky_tests=flaky_tests,
        stable_tests=stable_count,
        scan_duration=time.time() - start_time,
        runs_per_test=runs,
    )


def format_flaky_scan_report(result: FlakyScanResult) -> str:
    """Format flaky scan results as a human-readable report.

    Args:
        result: FlakyScanResult to format

    Returns:
        Formatted report string
    """
    lines = [
        "",
        "Flaky Test Scan Report",
        "=" * 60,
        "",
        f"Total tests scanned: {result.total_tests}",
        f"Runs per test:       {result.runs_per_test}",
        f"Scan duration:       {result.scan_duration:.1f}s",
        "",
        f"Stable tests:        {result.stable_tests}",
        f"Flaky tests found:   {len(result.flaky_tests)}",
        "",
    ]

    if result.flaky_tests:
        lines.append("Flaky Tests (sorted by flakiness):")
        lines.append("-" * 60)

        for test in result.flaky_tests:
            lines.append(f"\n  {test.nodeid}")
            lines.append(f"    Pass rate:    {test.pass_rate*100:.0f}% ({test.pass_count}/{test.total_runs})")
            lines.append(f"    Flakiness:    {test.flakiness_score:.0f}%")
            if test.failure_signatures:
                lines.append(f"    Errors:       {', '.join(test.failure_signatures)}")

        lines.append("")
        lines.append("-" * 60)
        lines.append(f"Run with --quarantine to add these to {QUARANTINE_FILE}")
    else:
        lines.append("No flaky tests detected! All tests are stable.")

    return "\n".join(lines)


def add_to_quarantine(tests: List[QuarantinedTest], path: Optional[Path] = None) -> int:
    """Add flaky tests to the quarantine file.

    Args:
        tests: List of tests to quarantine
        path: Path to quarantine file

    Returns:
        Number of tests added
    """
    qf = load_quarantine_file(path)

    added = 0
    for test in tests:
        if not qf.is_quarantined(test.nodeid):
            added += 1
        qf.add_test(test)

    save_quarantine_file(qf, path)
    return added


def remove_from_quarantine(nodeids: List[str], path: Optional[Path] = None) -> int:
    """Remove tests from quarantine.

    Args:
        nodeids: List of test nodeids to remove
        path: Path to quarantine file

    Returns:
        Number of tests removed
    """
    qf = load_quarantine_file(path)

    removed = 0
    for nodeid in nodeids:
        if qf.remove_test(nodeid):
            removed += 1

    save_quarantine_file(qf, path)
    return removed


def list_quarantined_tests(path: Optional[Path] = None) -> List[QuarantinedTest]:
    """List all quarantined tests.

    Args:
        path: Path to quarantine file

    Returns:
        List of quarantined tests
    """
    qf = load_quarantine_file(path)
    return list(qf.tests.values())


def generate_pytest_deselect_args(path: Optional[Path] = None) -> List[str]:
    """Generate pytest arguments to skip quarantined tests.

    Args:
        path: Path to quarantine file

    Returns:
        List of pytest arguments (e.g., ["--deselect", "test::nodeid", ...])
    """
    qf = load_quarantine_file(path)

    args = []
    for nodeid in qf.tests:
        args.extend(["--deselect", nodeid])

    return args
