"""
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 coverage import CoverageData, CoverageException

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[list[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) ---
    # We flatten the per-test coverage data ("bitmaps") into a
    # "long-form" list, which is perfect for database ingestion.
    # Structure: [{'file': str, 'line': int, 'context': str},...]
    coverage_list = []

    for filename in measured_files:
        # contexts_by_lineno() is the key API call.
        # It returns a dict: {line_num: [context1, context2]}
        # This is the "bitmap" data for this file.
        try:
            contexts_map = cov_data.contexts_by_lineno(filename)
        except Exception as e:
            console.print(f"[yellow]Warning: Could not process file {filename}: {e}[/yellow]")
            continue

        try:
            for line_num, contexts in contexts_map.items():
                for context_name in contexts:
                    # Accept empty contexts (global context) as valid
                    # Strip everything after | from context names for cleaner processing
                    clean_context = context_name.split("|")[0] if context_name else "global"
                    coverage_list.append(
                        {
                            "file": filename,
                            "line": line_num,
                            "from_line": None,
                            "to_line": None,
                            "context": clean_context,
                            "type": "line",
                        }
                    )
        except Exception as e:
            console.print(f"[red]Error processing contexts for {filename}: {e}[/red]")
            console.print(f"[dim]contexts_map = {contexts_map}[/dim]")
            continue

        # Also extract branch (arc) coverage if requested
        if with_branches:
            try:
                # For arcs, we need to query the database directly since there's no contexts_by_arcs API
                import sqlite3

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

                    # Single query to get all arcs for all measured files
                    placeholders = ", ".join(["?"] * len(measured_files))
                    arcs_data = cursor.execute(
                        f"""
                        SELECT f.path, a.fromno, a.tono, c.context
                        FROM file f
                        JOIN arc a ON f.id = a.file_id
                        JOIN context c ON a.context_id = c.id
                        WHERE f.path IN ({placeholders})
                    """,
                        list(measured_files),
                    ).fetchall()

                    for filename, from_line, to_line, context_name in arcs_data:
                        clean_context = context_name.split("|")[0] if context_name else "global"
                        coverage_list.append(
                            {
                                "file": filename,
                                "from_line": from_line,
                                "to_line": to_line,
                                "context": clean_context,
                                "type": "arc",
                            }
                        )
            except Exception as e:
                console.print(f"[yellow]Warning: Could not process arcs: {e}[/yellow]")
                continue

    if not coverage_list:
        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 {len(coverage_list)} coverage entries from {len({item['context'] for item in coverage_list})} unique test contexts."
    )

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

    return coverage_list, source_files, has_branches


def calculate_redundancy_analysis(
    coverage_list: list[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
    """
    console = Console()

    # --- 3. LOAD (Into DuckDB via PyArrow) ---
    console.print("Loading data into in-memory analytical database...")

    # Normalize data to ensure consistent schema
    if coverage_list:
        all_keys = set()
        for item in coverage_list:
            all_keys.update(item.keys())
        # Modify the list in place
        for i, item in enumerate(coverage_list):
            normalized_item = {key: item.get(key) for key in all_keys}
            coverage_list[i] = normalized_item

    import pyarrow as pa

    if not coverage_list:
        console.print("[yellow]No coverage data to analyze.[/yellow]")
        return []

    try:
        # pa.Table.from_pylist is highly efficient for this
        coverage_table = pa.Table.from_pylist(coverage_list)
    except Exception as e:
        print(f"Error converting data to PyArrow table: {e}", file=sys.stderr)
        return []

    import duckdb

    con = duckdb.connect(database=":memory:")

    # Register the PyArrow table as a virtual SQL table. Zero-copy.
    con.register("raw_coverage_data", coverage_table)

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

    # Get all unique test contexts
    test_contexts = con.execute(
        "SELECT DISTINCT context FROM raw_coverage_data WHERE context != 'global'"
    ).fetchall()
    test_contexts = [row[0] for row in test_contexts]

    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.")

    # Load full test coverage data (needed for verbose output)
    # Use a single grouped query to aggregate coverage data by context
    context_placeholders = ", ".join([f"'{ctx}'" for ctx in test_contexts])

    # Check if arc columns exist in the table schema
    table_columns = set(con.execute("DESCRIBE raw_coverage_data").fetchall())
    table_columns = {row[0] for row in table_columns}
    has_arc_columns = "from_line" in table_columns and "to_line" in table_columns

    # Check if arc data exists
    try:
        has_arcs = (
            con.execute(f"""
            SELECT 1 FROM raw_coverage_data
            WHERE context IN ({context_placeholders}) AND type = 'arc'
            LIMIT 1
        """).fetchone()
            is not None
        )
    except Exception:
        has_arcs = False

    if has_arc_columns and has_arcs:
        coverage_data = con.execute(f"""
            SELECT
                context,
                array_agg(struct_pack(file := file, line := line)) FILTER (WHERE type = 'line') as lines,
                array_agg(struct_pack(file := file, from_line := COALESCE(from_line, 0), to_line := COALESCE(to_line, 0))) FILTER (WHERE type = 'arc') as arcs
            FROM raw_coverage_data
            WHERE context IN ({context_placeholders})
            GROUP BY context
        """).fetchall()
    else:
        coverage_data = (
            con.execute(f"""
            SELECT
                context,
                array_agg(struct_pack(file := file, line := line)) FILTER (WHERE type = 'line') as lines,
                [] as arcs
            FROM raw_coverage_data
            WHERE context IN ({context_placeholders})
            GROUP BY context
        """).fetchall()
        )  # Process aggregated results into test_coverage and test_source_coverage dictionaries
    test_coverage = {}
    test_source_coverage = {}

    for row in coverage_data:
        context, lines_array, arcs_array = row

        # Convert arrays of structs to sets of tuples
        all_lines = {(item["file"], item["line"]) for item in lines_array} if lines_array else set()
        all_arcs = (
            {(item["file"], item["from_line"], item["to_line"]) for item in arcs_array}
            if arcs_array
            else set()
        )

        # Filter for source files
        source_lines = {(file, line) for file, line in all_lines if file in source_files}
        source_arcs = {(file, fr, to) for file, fr, to in all_arcs if file in source_files}

        test_coverage[context] = {
            "lines": all_lines,
            "arcs": all_arcs,
        }

        test_source_coverage[context] = {
            "lines": source_lines,
            "arcs": source_arcs,
        }

    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...")

    union_coverage = {}
    for test_context in test_contexts:
        other_tests = [t for t in test_contexts if t != test_context]
        if not other_tests:
            union_coverage[test_context] = {"lines": set(), "arcs": set()}
            continue

        union_lines = set()
        union_arcs = set()
        for other_test in other_tests:
            union_lines.update(test_source_coverage[other_test]["lines"])
            if with_branches:
                union_arcs.update(test_source_coverage[other_test]["arcs"])

        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 overlap percentages
        lines_overlap = len(test_lines & union_lines) / len(test_lines) * 100 if test_lines else 0.0
        arcs_overlap = len(test_arcs & union_arcs) / len(test_arcs) * 100 if test_arcs else 0.0

        # A test is redundant if its coverage is completely covered by other tests
        lines_redundant = not test_lines or test_lines.issubset(union_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 = len(test_coverage[test_context]["lines"])
        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
    total_lines_covered = len(
        {(f, line) for result in test_results for f, line in result["source_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
                lines_by_file = {}
                for file, line in result["source_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 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 = {
                (f, line_num)
                for f, line_num in test_coverage[test_context]["lines"]
                if f not in source_files
            }
            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, uses DuckDB to perform a
    set-based analysis, and prints redundant tests.

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

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

    # Build test_coverage dict for verbose output (needed by formatting function)
    import duckdb
    import pyarrow as pa

    coverage_table = pa.Table.from_pylist(coverage_list)
    con = duckdb.connect(database=":memory:")
    con.register("raw_coverage_data", coverage_table)

    test_contexts = [result["test"] for result in test_results]
    test_coverage = {}
    for test_context in test_contexts:
        lines = con.execute(f"""
            SELECT file, line
            FROM raw_coverage_data
            WHERE context = '{test_context}' AND type = 'line'
        """).fetchall()
        arcs = con.execute(f"""
            SELECT file, from_line, to_line
            FROM raw_coverage_data
            WHERE context = '{test_context}' AND type = 'arc'
        """).fetchall()

        test_coverage[test_context] = {
            "lines": {(row[0], row[1]) for row in lines},
            "arcs": {(row[0], row[1], row[2]) for row in arcs},
        }

    # 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)
