"""
Analyze test redundancy using coverage data.

This script reads the .coverage database generated by pytest-cov and performs
redundancy analysis to identify tests that don't contribute unique coverage.

The analysis uses a three-pass algorithm:
1. Extract coverage data from .coverage database
2. Calculate union coverage of all other tests for each test
3. Determine redundancy based on overlap percentages

Usage:
    uvx recov [options]

Examples:
    # Basic analysis (lines only)
    uvx recov

    # Verbose mode with detailed coverage
    uvx recov -v

    # Require branch coverage analysis
    uvx recov --with-branches

    # Verbose with branches
    uvx recov -v --with-branches

Prerequisites:
    - Run pytest with coverage: pytest --cov=src --cov-context=test --cov-branch --cov-append
    - Ensure .coverage file exists in current directory
"""

import argparse
import sys
from collections import defaultdict

from coverage import CoverageData, CoverageException
from coverage.numbits import (
    nums_to_numbits,
    numbits_union,
    numbits_intersection,
    numbits_to_nums,
)

try:
    from rich.console import Console
    from rich.table import Table
except ImportError:
    print("Error: rich library is required. Install with: uv add rich", file=sys.stderr)
    sys.exit(1)


def extract_coverage_data(
    with_branches: bool = False,
) -> tuple[dict, set[str], bool]:
    """
    Extract coverage data from .coverage file and prepare it for analysis.

    Returns:
        tuple: (coverage_list, source_files, has_branches)
    """
    console = Console()

    # --- 1. EXTRACT (Build the raw data for our "bitmaps") ---
    console.print("[bold blue]Reading .coverage database...[/bold blue]")
    try:
        cov_data = CoverageData()  # no_disk=True
        cov_data.read()
    except CoverageException as e:
        console.print(f"[red]Error reading .coverage file: {e}[/red]")
        console.print(
            "[yellow]Ensure a .coverage file was generated by pytest-cov.[/yellow]"
        )
        return {}, set(), False

    measured_files = cov_data.measured_files()
    if not measured_files:
        console.print("[red]No measured files found in .coverage data.[/red]")
        return {}, set(), False

    console.print(
        f"Found {len(measured_files)} measured files. Extracting line and branch contexts..."
    )

    # Check if branch coverage is available
    has_branches = cov_data.has_arcs()
    if with_branches:
        if not has_branches:
            console.print(
                "[red]✗ Branch coverage required but not available in .coverage file[/red]"
            )
            console.print(
                "[yellow]Use --cov-branch when running pytest to collect branch coverage[/yellow]"
            )
            return {}, set(), False
        console.print(
            "[green]✓ Branch coverage data available and will be analyzed[/green]"
        )
    elif has_branches:
        console.print("[green]✓ Branch coverage data available[/green]")
    else:
        console.print(
            "[dim]ℹ Line coverage only (use --with-branches to require branch analysis)[/dim]"
        )

    # --- 2. TRANSFORM (Convert data to a queryable format) ---
    # Build test coverage data structure directly
    # Structure: {context: {"lines": {file: numbits}, "arcs": set of (file, fromno, tono)}}
    test_coverage = defaultdict(lambda: {"lines": {}, "arcs": set()})

    try:
        import sqlite3

        with sqlite3.connect(cov_data.data_filename()) as conn:
            cursor = conn.cursor()

            # Bulk queries for all measured files
            # Query for line_bits data
            lines_query = """
                SELECT f.path, l.numbits, c.context
                FROM line_bits l
                JOIN context c ON l.context_id = c.id
                JOIN file f ON l.file_id = f.id
            """
            lines_data = cursor.execute(lines_query).fetchall()

            # Query for arcs data
            arcs_query = """
                SELECT f.path, arc.fromno, arc.tono, c.context
                FROM arc
                JOIN context c ON arc.context_id = c.id
                JOIN file f ON arc.file_id = f.id
            """
            arcs_data = cursor.execute(arcs_query).fetchall()

    except Exception as e:
        console.print(f"[yellow]Warning: Could not process coverage data: {e}[/yellow]")
        return {}, set(), False

    # Process lines data
    for filename, numbits, context_name in lines_data:
        clean_context = context_name.split("|")[0] if context_name else "global"
        existing = test_coverage[clean_context]["lines"].get(filename, b"")
        if existing:
            test_coverage[clean_context]["lines"][filename] = numbits_union(
                existing, numbits
            )
        else:
            test_coverage[clean_context]["lines"][filename] = numbits

    # Process arcs data
    for filename, fromno, tono, context_name in arcs_data:
        clean_context = context_name.split("|")[0] if context_name else "global"
        test_coverage[clean_context]["arcs"].add((filename, fromno, tono))

    # Infer lines from arcs for files that don't have line_bits
    files_with_line_bits = {filename for filename, _, _ in lines_data}
    for filename in measured_files:
        if filename not in files_with_line_bits:
            # Infer lines from arcs
            lines_by_context = defaultdict(set)
            arcs_for_file = []
            for context in test_coverage:
                for arc in test_coverage[context]["arcs"]:
                    if arc[0] == filename:
                        arcs_for_file.append((context, arc[1], arc[2]))
            for context, fromno, tono in arcs_for_file:
                if fromno > 0:
                    lines_by_context[context].add(fromno)
                if tono > 0:
                    lines_by_context[context].add(tono)
            for context, lines in lines_by_context.items():
                if lines:
                    numbits = nums_to_numbits(sorted(lines))
                    test_coverage[context]["lines"][filename] = numbits

    if not test_coverage:
        console.print("[red]No context data available in .coverage file.[/red]")
        console.print("[yellow]Ensure pytest was run with --cov-context=test[/yellow]")
        return {}, set(), False

    console.print(f"Extracted coverage for {len(test_coverage)} unique test contexts.")

    # Use all coverage data for analysis, but track source vs test coverage separately
    source_files = {f for f in measured_files}
    console.print(
        f"Identified {len(source_files)} source files for redundancy analysis."
    )

    return test_coverage, source_files, has_branches


def calculate_redundancy_analysis(
    test_coverage: dict, source_files: set[str], with_branches: bool
) -> list[dict]:
    """
    Perform the three-pass redundancy analysis on coverage data.

    Returns:
        list: test_results with redundancy analysis data
    """
    import decimal

    console = Console()

    # --- 3. ANALYZE (Three-pass redundancy detection) ---

    # Get all unique test contexts
    test_contexts = [ctx for ctx in test_coverage if ctx != "global"]

    if not test_contexts:
        console.print("[yellow]No test-specific coverage contexts found.[/yellow]")
        console.print("[yellow]Ensure pytest was run with --cov-context=test[/yellow]")
        return []

    console.print(f"Found {len(test_contexts)} test contexts to analyze.")

    # Build test_source_coverage (filtered to source files)
    test_source_coverage = {}
    for ctx in test_contexts:
        test_source_coverage[ctx] = {
            "lines": {
                file: numbits
                for file, numbits in test_coverage[ctx]["lines"].items()
                if file in source_files
            },
            "arcs": {
                (file, fr, to)
                for file, fr, to in test_coverage[ctx]["arcs"]
                if file in source_files
            },
        }

    console.print(f"Calculated coverage for {len(test_source_coverage)} tests.")

    # PASS 2: For each test, calculate union coverage of all other tests except that one
    console.print("Pass 2: Calculating union coverage of other tests...")

    # Use dynamic programming to precompute prefix and suffix unions for efficient computation
    sorted_contexts = sorted(test_contexts)
    n = len(sorted_contexts)

    # Prefix unions: prefix_numbits[file][i] = union of test_source_coverage[sorted_contexts[0]] to sorted_contexts[i-1] for that file
    prefix_numbits = defaultdict(lambda: [b""] * (n + 1))
    prefix_arcs = [set()] * (n + 1) if with_branches else None

    for i in range(1, n + 1):
        current_lines = test_source_coverage[sorted_contexts[i - 1]]["lines"]
        for file, numbits in current_lines.items():
            prefix_numbits[file][i] = numbits_union(
                prefix_numbits[file][i - 1], numbits
            )
        if with_branches:
            prev_union_arcs = prefix_arcs[i - 1]
            current_arcs = test_source_coverage[sorted_contexts[i - 1]]["arcs"]
            prefix_arcs[i] = prev_union_arcs | current_arcs

    # Suffix unions: suffix_numbits[file][i] = union of test_source_coverage[sorted_contexts[i]] to sorted_contexts[n-1] for that file
    suffix_numbits = defaultdict(lambda: [b""] * (n + 1))
    suffix_arcs = [set()] * (n + 1) if with_branches else None
    for file in prefix_numbits:
        suffix_numbits[file][n] = b""
    if with_branches:
        suffix_arcs[n] = set()

    for i in range(n - 1, -1, -1):
        current_lines = test_source_coverage[sorted_contexts[i]]["lines"]
        for file, numbits in current_lines.items():
            suffix_numbits[file][i] = numbits_union(
                numbits, suffix_numbits[file][i + 1]
            )
        if with_branches:
            next_union_arcs = suffix_arcs[i + 1]
            current_arcs = test_source_coverage[sorted_contexts[i]]["arcs"]
            suffix_arcs[i] = current_arcs | next_union_arcs

    # Now compute union_coverage using prefix and suffix
    union_coverage = {}
    for idx, test_context in enumerate(sorted_contexts):
        union_lines = {}
        for file in prefix_numbits:
            union_lines[file] = numbits_union(
                prefix_numbits[file][idx], suffix_numbits[file][idx + 1]
            )
        union_arcs = (
            (prefix_arcs[idx] | suffix_arcs[idx + 1]) if with_branches else set()
        )
        union_coverage[test_context] = {
            "lines": union_lines,
            "arcs": union_arcs,
        }

    console.print("Calculated union coverage for all tests.")

    # PASS 3: Calculate overlap (%) of coverage
    console.print("Pass 3: Calculating redundancy overlap percentages...")

    test_results = []

    for test_context in test_contexts:
        test_lines = test_source_coverage[test_context]["lines"]
        test_arcs = test_source_coverage[test_context]["arcs"]
        union_lines = union_coverage[test_context]["lines"]
        union_arcs = union_coverage[test_context]["arcs"]

        # Calculate total lines covered by this test
        total_test_lines = sum(
            len(numbits_to_nums(numbits)) for numbits in test_lines.values()
        )

        # Calculate overlap for lines
        if test_lines:
            overlap_lines = 0
            for file in test_lines:
                if file in union_lines:
                    intersection = numbits_intersection(
                        test_lines[file], union_lines[file]
                    )
                    overlap_lines += len(numbits_to_nums(intersection))
            lines_overlap = (
                decimal.Decimal(overlap_lines)
                / decimal.Decimal(total_test_lines)
                * decimal.Decimal(100)
            ).quantize(decimal.Decimal("0.01"))
        else:
            lines_overlap = decimal.Decimal("0.00")

        # Calculate overlap for arcs (unchanged)
        if test_arcs:
            overlap_count = decimal.Decimal(len(test_arcs & union_arcs))
            total_count = decimal.Decimal(len(test_arcs))
            arcs_overlap = (
                overlap_count / total_count * decimal.Decimal(100)
            ).quantize(decimal.Decimal("0.01"))
        else:
            arcs_overlap = decimal.Decimal("0.00")

        # A test is redundant if its coverage is completely covered by other tests
        lines_redundant = not test_lines or all(
            numbits_intersection(test_lines[file], union_lines.get(file, b""))
            == test_lines[file]
            for file in test_lines
        )
        arcs_redundant = not test_arcs or test_arcs.issubset(union_arcs)

        # If branches are not being analyzed, only consider line redundancy
        is_redundant = lines_redundant and (arcs_redundant if with_branches else True)

        # Total covered items (lines + arcs, including test files)
        total_lines = sum(
            len(numbits_to_nums(numbits))
            for numbits in test_coverage[test_context]["lines"].values()
        )
        total_arcs = len(test_coverage[test_context]["arcs"])
        total_covered_items = total_lines + total_arcs

        test_results.append(
            {
                "test": test_context,
                "lines_overlap": lines_overlap,
                "arcs_overlap": arcs_overlap,
                "is_redundant": is_redundant,
                "total_covered_items": total_covered_items,
                "source_lines": test_lines,
                "source_arcs": test_arcs,
            }
        )

    console.print("Calculated redundancy analysis.")

    return test_results


def format_and_display_results(
    test_results: list[dict],
    test_coverage: dict,
    source_files: set[str],
    with_branches: bool,
    verbose: bool,
) -> None:
    """
    Format and display the redundancy analysis results.
    """
    console = Console()

    # --- 5. REPORT (Results) ---
    console.print("\n[bold blue]=== Test Redundancy Analysis Results ===[/bold blue]")

    # Sort test results by redundancy status (redundant first) then by total covered items
    test_results.sort(key=lambda x: (not x["is_redundant"], -x["total_covered_items"]))

    # Show all tests in a table
    analysis_type = "Lines + Branches" if with_branches else "Lines Only"
    table = Table(title=f"All Tests (Sorted by Redundancy Status) - {analysis_type}")
    table.add_column("Test", style="cyan")
    table.add_column("Covered Items", style="magenta")
    table.add_column("Lines Overlap %", style="yellow")
    if with_branches:
        table.add_column("Arcs Overlap %", style="yellow")
    table.add_column(
        "Status",
        style="red" if any(r["is_redundant"] for r in test_results) else "green",
    )

    for result in test_results:
        status = "REDUNDANT" if result["is_redundant"] else "UNIQUE"
        row_data = [
            result["test"],
            str(result["total_covered_items"]),
            f"{result['lines_overlap']:.1f}%" if result["source_lines"] else "N/A",
        ]
        if with_branches:
            row_data.append(
                f"{result['arcs_overlap']:.1f}%" if result["source_arcs"] else "N/A"
            )
        row_data.append(status)
        table.add_row(*row_data)

    console.print(table)

    # Summary of redundant tests
    redundant_count = sum(1 for r in test_results if r["is_redundant"])
    analysis_basis = "lines and branches" if with_branches else "lines only"
    if redundant_count > 0:
        console.print(
            f"\n[red]⚠ Found {redundant_count} redundant test(s) (coverage completely covered by other tests based on {analysis_basis}):[/red]"
        )
        for result in test_results:
            if result["is_redundant"]:
                console.print(f"  • {result['test']}")
    else:
        console.print("\n[green]✓ No redundant tests found![/green]")
        console.print(
            f"All tests contribute unique coverage that isn't fully covered by other tests (analyzed {analysis_basis})."
        )

    # Coverage summary
    all_lines = set()
    for result in test_results:
        for file, numbits in result["source_lines"].items():
            all_lines.update((file, line) for line in numbits_to_nums(numbits))
    total_lines_covered = len(all_lines)
    total_arcs_covered = (
        len(
            {
                (f, fr, to)
                for result in test_results
                for f, fr, to in result["source_arcs"]
            }
        )
        if with_branches
        else 0
    )
    console.print("\n[bold]Coverage Summary:[/bold]")
    console.print(f"• Total tests analyzed: {len(test_results)}")
    console.print(f"• Total lines covered: {total_lines_covered}")
    if with_branches:
        console.print(f"• Total arcs covered: {total_arcs_covered}")
    console.print(f"• Redundant tests: {redundant_count}")

    if verbose:
        console.print("\n[bold cyan]=== Detailed Coverage per Test ===[/bold cyan]")
        for result in test_results:
            test_context = result["test"]
            console.print(f"\n[bold]{test_context}[/bold]")
            console.print(
                f"  Status: {'REDUNDANT' if result['is_redundant'] else 'UNIQUE'}"
            )
            console.print(f"  Lines overlap: {result['lines_overlap']:.1f}%")
            if with_branches:
                console.print(f"  Arcs overlap: {result['arcs_overlap']:.1f}%")
            console.print(f"  Total covered items: {result['total_covered_items']}")

            # Show source file coverage
            if result["source_lines"]:
                console.print("  [blue]Source files (lines):[/blue]")
                # Group by file
                for file, numbits in result["source_lines"].items():
                    lines = sorted(numbits_to_nums(numbits))
                    # Group consecutive lines
                    ranges = []
                    if lines:
                        start = lines[0]
                        prev = start
                        for num in lines[1:]:
                            if num == prev + 1:
                                prev = num
                            else:
                                if start == prev:
                                    ranges.append(str(start))
                                else:
                                    ranges.append(f"{start}-{prev}")
                                start = num
                                prev = num
                        if start == prev:
                            ranges.append(str(start))
                        else:
                            ranges.append(f"{start}-{prev}")
                    console.print(f"    {file}: {', '.join(ranges)}")

            if with_branches and result["source_arcs"]:
                console.print("  [blue]Source files (branches):[/blue]")
                arcs_by_file = {}
                for file, from_line, to_line in result["source_arcs"]:
                    if file not in arcs_by_file:
                        arcs_by_file[file] = []
                    arcs_by_file[file].append((from_line, to_line))

                for file, arcs in arcs_by_file.items():
                    arcs.sort()
                    arc_strs = [f"{fr}→{to}" for fr, to in arcs]
                    console.print(f"    {file}: {', '.join(arc_strs)}")

            # Show test/other file coverage
            test_lines = set()
            for f, numbits in test_coverage[test_context]["lines"].items():
                if f not in source_files:
                    test_lines.update(
                        (f, line_num) for line_num in numbits_to_nums(numbits)
                    )
            test_arcs = (
                {
                    (f, fr, to)
                    for f, fr, to in test_coverage[test_context]["arcs"]
                    if f not in source_files
                }
                if with_branches
                else set()
            )

            if test_lines:
                console.print("  [blue]Test/other files (lines):[/blue]")
                lines_by_file = {}
                for file, line in test_lines:
                    if file not in lines_by_file:
                        lines_by_file[file] = []
                    lines_by_file[file].append(line)

                for file, lines in lines_by_file.items():
                    lines.sort()
                    # Group consecutive lines
                    ranges = []
                    start = lines[0]
                    prev = start
                    for num in lines[1:]:
                        if num == prev + 1:
                            prev = num
                        else:
                            if start == prev:
                                ranges.append(str(start))
                            else:
                                ranges.append(f"{start}-{prev}")
                            start = num
                            prev = num
                    if start == prev:
                        ranges.append(str(start))
                    else:
                        ranges.append(f"{start}-{prev}")
                    console.print(f"    {file}: {', '.join(ranges)}")

            if with_branches and test_arcs:
                console.print("  [blue]Test/other files (branches):[/blue]")
                arcs_by_file = {}
                for file, from_line, to_line in test_arcs:
                    if file not in arcs_by_file:
                        arcs_by_file[file] = []
                    arcs_by_file[file].append((from_line, to_line))

                for file, arcs in arcs_by_file.items():
                    arcs.sort()
                    arc_strs = [f"{fr}→{to}" for fr, to in arcs]
                    console.print(f"    {file}: {', '.join(arc_strs)}")


def analyze_coverage_for_redundancy(
    verbose: bool = False, with_branches: bool = False
) -> None:
    """
    Reads the .coverage database, performs redundancy analysis, and prints redundant tests.

    This is the implementation of the "bitmap" and "NlogN" set-difference algorithm.
    """
    # Extract coverage data
    test_coverage, source_files, has_branches = extract_coverage_data(with_branches)
    if not test_coverage:
        return

    # Perform redundancy analysis
    test_results = calculate_redundancy_analysis(
        test_coverage, source_files, with_branches
    )
    if not test_results:
        return

    # Format and display results
    format_and_display_results(
        test_results, test_coverage, source_files, with_branches, verbose
    )


def main() -> None:
    """Entry point for uv script execution."""
    parser = argparse.ArgumentParser(
        description="Analyze test redundancy using coverage data",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Local development
  uvx recov

  # Standalone script (anywhere)
  uvx recov --with-branches
  uvx recov -v --with-branches
""",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        help="Show detailed coverage information for each test",
    )
    parser.add_argument(
        "--with-branches",
        action="store_true",
        help="Require branch coverage analysis (fails if branches not available)",
    )

    args = parser.parse_args()
    analyze_coverage_for_redundancy(
        verbose=args.verbose, with_branches=args.with_branches
    )
