"""Spectrum-Based Fault Localization (SBFL).

Implements statistical fault localization techniques from academic research
(Google's De-Flake paper, etc.) to rank code locations by suspiciousness.

For each code location, we calculate how often it's executed in:
- Failing tests (ef, ep)
- Passing tests (nf, np)

Then apply formulas like Ochiai or Tarantula to rank suspiciousness.

Usage:
    sbfl = SBFLAnalyzer()
    suspicious = sbfl.analyze(coverage_data, test_results)

    for loc in suspicious[:10]:
        print(f"{loc.file}:{loc.line} - score: {loc.score:.3f}")
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from enum import Enum
from typing import Dict, List, Optional, Set, Tuple


class SBFLFormula(Enum):
    """Available SBFL formulas."""
    OCHIAI = "ochiai"  # Best overall performance
    TARANTULA = "tarantula"  # Classic formula
    DStar = "dstar"  # Good for multiple faults
    JACCARD = "jaccard"
    OP2 = "op2"  # Recent, good performance


@dataclass
class LineStats:
    """Statistics for a single line of code."""
    ef: int  # Executed in failing tests
    ep: int  # Executed in passing tests
    nf: int  # Not executed in failing tests
    np: int  # Not executed in passing tests
    total_failed: int
    total_passed: int


@dataclass
class SuspiciousLocation:
    """A suspicious code location with its score."""
    file: str
    line: int
    function: Optional[str]
    score: float
    formula: str
    stats: LineStats
    rank: int = 0

    @property
    def location_str(self) -> str:
        """Format as file:line string."""
        return f"{self.file}:{self.line}"


@dataclass
class TestResult:
    """Result of a single test run."""
    test_name: str
    passed: bool
    covered_lines: Dict[str, Set[int]]  # file -> set of covered line numbers


@dataclass
class CoverageData:
    """Coverage data for SBFL analysis."""
    # Map of (file, line) -> set of tests that executed it
    line_to_tests: Dict[Tuple[str, int], Set[str]]
    # All test results
    test_results: Dict[str, bool]  # test_name -> passed
    # Optional: function names for lines
    line_to_function: Dict[Tuple[str, int], str]


class SBFLAnalyzer:
    """Spectrum-Based Fault Localization analyzer."""

    def __init__(self, formula: SBFLFormula = SBFLFormula.OCHIAI):
        """Initialize SBFL analyzer.

        Args:
            formula: The suspiciousness formula to use.
        """
        self.formula = formula

    def analyze(
        self,
        coverage: CoverageData,
        max_results: int = 100,
    ) -> List[SuspiciousLocation]:
        """Analyze coverage data to find suspicious locations.

        Args:
            coverage: Coverage data with test execution info.
            max_results: Maximum number of results to return.

        Returns:
            List of suspicious locations, sorted by score (highest first).
        """
        # Count passing and failing tests
        total_passed = sum(1 for p in coverage.test_results.values() if p)
        total_failed = sum(1 for p in coverage.test_results.values() if not p)

        if total_failed == 0:
            return []  # No failures, nothing to localize

        # Calculate stats for each line
        line_stats: Dict[Tuple[str, int], LineStats] = {}

        for (file, line), tests in coverage.line_to_tests.items():
            ef = 0  # Executed in failing
            ep = 0  # Executed in passing

            for test in tests:
                if test in coverage.test_results:
                    if coverage.test_results[test]:
                        ep += 1
                    else:
                        ef += 1

            nf = total_failed - ef
            np = total_passed - ep

            line_stats[(file, line)] = LineStats(
                ef=ef, ep=ep, nf=nf, np=np,
                total_failed=total_failed,
                total_passed=total_passed
            )

        # Calculate suspiciousness scores
        suspicious: List[SuspiciousLocation] = []

        for (file, line), stats in line_stats.items():
            score = self._calculate_score(stats)

            if score > 0:  # Only include lines with some suspiciousness
                function = coverage.line_to_function.get((file, line))
                suspicious.append(SuspiciousLocation(
                    file=file,
                    line=line,
                    function=function,
                    score=score,
                    formula=self.formula.value,
                    stats=stats,
                ))

        # Sort by score (highest first) and assign ranks
        suspicious.sort(key=lambda s: s.score, reverse=True)
        for i, loc in enumerate(suspicious[:max_results]):
            loc.rank = i + 1

        return suspicious[:max_results]

    def _calculate_score(self, stats: LineStats) -> float:
        """Calculate suspiciousness score using selected formula."""
        if self.formula == SBFLFormula.OCHIAI:
            return self._ochiai(stats)
        elif self.formula == SBFLFormula.TARANTULA:
            return self._tarantula(stats)
        elif self.formula == SBFLFormula.DStar:
            return self._dstar(stats)
        elif self.formula == SBFLFormula.JACCARD:
            return self._jaccard(stats)
        elif self.formula == SBFLFormula.OP2:
            return self._op2(stats)
        else:
            return self._ochiai(stats)

    def _ochiai(self, stats: LineStats) -> float:
        """Ochiai formula - best overall SBFL performance.

        score = ef / sqrt(total_failed * (ef + ep))

        This is based on the Ochiai similarity coefficient from genetics.
        """
        if stats.total_failed == 0 or (stats.ef + stats.ep) == 0:
            return 0.0
        denominator = math.sqrt(stats.total_failed * (stats.ef + stats.ep))
        if denominator == 0:
            return 0.0
        return stats.ef / denominator

    def _tarantula(self, stats: LineStats) -> float:
        """Tarantula formula - classic SBFL.

        score = (ef/total_failed) / (ef/total_failed + ep/total_passed)
        """
        if stats.total_failed == 0 or stats.total_passed == 0:
            return 0.0

        fail_ratio = stats.ef / stats.total_failed if stats.total_failed > 0 else 0
        pass_ratio = stats.ep / stats.total_passed if stats.total_passed > 0 else 0

        denominator = fail_ratio + pass_ratio
        if denominator == 0:
            return 0.0
        return fail_ratio / denominator

    def _dstar(self, stats: LineStats, star: int = 2) -> float:
        """D* formula - good for multiple faults.

        score = ef^* / (ep + nf)

        Default star=2 (D*2) is commonly used.
        Returns normalized score in [0, 1] for cross-formula comparability.
        """
        denominator = stats.ep + stats.nf
        if denominator == 0:
            # All failing tests execute this line, no passing tests do —
            # maximally suspicious. Return 1.0 instead of infinity.
            return 1.0 if stats.ef > 0 else 0.0
        raw = (stats.ef ** star) / denominator
        # Normalize: cap at 1.0 for comparability with Ochiai/Tarantula/Jaccard
        return min(raw, 1.0)

    def _jaccard(self, stats: LineStats) -> float:
        """Jaccard formula.

        score = ef / (ef + nf + ep)
        """
        denominator = stats.ef + stats.nf + stats.ep
        if denominator == 0:
            return 0.0
        return stats.ef / denominator

    def _op2(self, stats: LineStats) -> float:
        """OP2 formula - recent, good performance.

        score = ef - ep / (ep + np + 1)

        Returns normalized score in [0, 1] for cross-formula comparability.
        """
        raw = stats.ef - (stats.ep / (stats.ep + stats.np + 1))
        # Normalize: OP2 can range from -1 to total_failed.
        # Map to [0, 1] where 0 = not suspicious, 1 = maximally suspicious.
        if stats.total_failed == 0:
            return 0.0
        return max(0.0, min(raw / stats.total_failed, 1.0))

    @classmethod
    def from_test_results(
        cls,
        test_results: List[TestResult],
        formula: SBFLFormula = SBFLFormula.OCHIAI,
    ) -> SBFLAnalyzer:
        """Create analyzer and coverage data from test results.

        Args:
            test_results: List of test results with coverage.
            formula: SBFL formula to use.

        Returns:
            Configured SBFLAnalyzer.
        """
        analyzer = cls(formula)

        # Build coverage data
        line_to_tests: Dict[Tuple[str, int], Set[str]] = {}
        test_passed: Dict[str, bool] = {}

        for result in test_results:
            test_passed[result.test_name] = result.passed

            for file, lines in result.covered_lines.items():
                for line in lines:
                    key = (file, line)
                    if key not in line_to_tests:
                        line_to_tests[key] = set()
                    line_to_tests[key].add(result.test_name)

        coverage = CoverageData(
            line_to_tests=line_to_tests,
            test_results=test_passed,
            line_to_function={},
        )

        return analyzer

    def analyze_from_results(
        self,
        test_results: List[TestResult],
        max_results: int = 100,
    ) -> List[SuspiciousLocation]:
        """Convenience method to analyze directly from test results."""
        # Build coverage data
        line_to_tests: Dict[Tuple[str, int], Set[str]] = {}
        test_passed: Dict[str, bool] = {}

        for result in test_results:
            test_passed[result.test_name] = result.passed

            for file, lines in result.covered_lines.items():
                for line in lines:
                    key = (file, line)
                    if key not in line_to_tests:
                        line_to_tests[key] = set()
                    line_to_tests[key].add(result.test_name)

        coverage = CoverageData(
            line_to_tests=line_to_tests,
            test_results=test_passed,
            line_to_function={},
        )

        return self.analyze(coverage, max_results)


def rank_suspicious_locations(
    passing_run_coverage: Dict[str, Set[int]],
    failing_run_coverage: Dict[str, Set[int]],
    formula: SBFLFormula = SBFLFormula.OCHIAI,
) -> List[SuspiciousLocation]:
    """Quick helper to rank locations from two runs (pass/fail).

    This is a simplified interface for the common case of comparing
    a passing run to a failing run.

    Args:
        passing_run_coverage: Coverage from passing run {file: {lines}}
        failing_run_coverage: Coverage from failing run {file: {lines}}
        formula: SBFL formula to use

    Returns:
        Ranked list of suspicious locations
    """
    test_results = [
        TestResult(
            test_name="passing_run",
            passed=True,
            covered_lines=passing_run_coverage,
        ),
        TestResult(
            test_name="failing_run",
            passed=False,
            covered_lines=failing_run_coverage,
        ),
    ]

    analyzer = SBFLAnalyzer(formula)
    return analyzer.analyze_from_results(test_results)
