"""Coverage integration for Python test runs.

Provides data types and helpers for loading, parsing, and aggregating
code-coverage information produced by coverage.py.
"""

from __future__ import annotations

import ast
import enum
import json
import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal

import attrs

from ..runtime._codeflash_wrap_decorator import get_run_tmp_file

if TYPE_CHECKING:
    from collections.abc import Collection

    from .._model import FunctionParent
    from ..context.models import CodeOptimizationContext

log = logging.getLogger(__name__)


class CoverageStatus(enum.Enum):
    """Status of a coverage-data lookup."""

    NOT_FOUND = "Coverage Data Not Found"
    PARSED_SUCCESSFULLY = "Parsed Successfully"


@attrs.frozen
class FunctionCoverage:
    """Coverage metrics for a single function."""

    name: str
    coverage: float
    executed_lines: list[int]
    unexecuted_lines: list[int]
    executed_branches: list[list[int]]
    unexecuted_branches: list[list[int]]


@attrs.frozen
class CoverageData:
    """Aggregated coverage data for an optimization target."""

    file_path: Path
    coverage: float
    function_name: str
    functions_being_tested: list[str]
    graph: dict[str, dict[str, Collection[object]]]
    code_context: CodeOptimizationContext
    main_func_coverage: FunctionCoverage
    dependent_func_coverage: FunctionCoverage | None
    status: CoverageStatus


def create_empty_coverage_data(
    file_path: Path,
    function_name: str,
    code_context: CodeOptimizationContext,
) -> CoverageData:
    """Create an empty :class:`CoverageData` with zero coverage."""
    return CoverageData(
        file_path=file_path,
        coverage=0.0,
        function_name=function_name,
        functions_being_tested=[function_name],
        graph={
            function_name: {
                "executed_lines": set(),
                "unexecuted_lines": set(),
                "executed_branches": [],
                "unexecuted_branches": [],
            }
        },
        code_context=code_context,
        main_func_coverage=FunctionCoverage(
            name=function_name,
            coverage=0.0,
            executed_lines=[],
            unexecuted_lines=[],
            executed_branches=[],
            unexecuted_branches=[],
        ),
        dependent_func_coverage=None,
        status=CoverageStatus.NOT_FOUND,
    )


def build_coverage_message(coverage_data: CoverageData) -> str:
    """Build a human-readable coverage summary."""
    if coverage_data.status == CoverageStatus.NOT_FOUND:
        return f"No coverage data found for {coverage_data.function_name}"
    return f"{coverage_data.coverage:.1f}%"


def extract_dependent_function(
    main_function: str,
    code_context: CodeOptimizationContext,
) -> str | Literal[False]:
    """Extract the single dependent function from *code_context* excluding *main_function*."""
    dependent_functions: set[str] = set()

    # Compare using bare name since AST extracts bare function names
    bare_main = (
        main_function.rsplit(".", 1)[-1]
        if "." in main_function
        else main_function
    )

    for code_string in code_context.testgen_context.code_strings:
        # Quick heuristic: skip parsing entirely if there is no 'def' token
        if "def" not in code_string.code:
            continue

        ast_tree = ast.parse(code_string.code)
        for node in ast_tree.body:
            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                name = node.name
                if name == bare_main:
                    continue
                dependent_functions.add(name)
                # Early exit when more than one dependent found
                if len(dependent_functions) > 1:
                    return False

    if not dependent_functions:
        return False

    if len(dependent_functions) != 1:
        return False

    return build_fully_qualified_name(dependent_functions.pop(), code_context)


def build_fully_qualified_name(
    function_name: str,
    code_context: CodeOptimizationContext,
) -> str:
    """Build a fully-qualified name for *function_name* using *code_context*."""
    # If the name is already qualified (contains a dot), return as-is
    if "." in function_name:
        return function_name
    full_name = function_name
    for obj_name, parents in code_context.preexisting_objects:
        if obj_name == function_name:
            parent: FunctionParent
            for parent in parents:
                if parent.type == "ClassDef":
                    full_name = f"{parent.name}.{full_name}"
            break
    return full_name


def generate_candidates(source_code_path: Path) -> set[str]:
    """Generate all possible candidates for coverage-data path matching."""
    candidates: set[str] = set()
    name = source_code_path.name
    candidates.add(name)

    parts = source_code_path.parts
    n = len(parts)

    last_added = name
    for i in range(n - 2, 0, -1):
        candidate_path = f"{parts[i]}/{last_added}"
        candidates.add(candidate_path)
        last_added = candidate_path

    candidates.add(source_code_path.as_posix())
    return candidates


def prepare_coverage_files() -> tuple[Path, Path]:
    """Prepare coverage configuration and output files."""
    coverage_database_file = get_run_tmp_file(Path(".coverage"))
    coveragercfile = get_run_tmp_file(Path(".coveragerc"))
    coveragerc_content = (
        f"[run]\n branch = True\ndata_file={coverage_database_file}\n"
    )
    coveragercfile.write_text(coveragerc_content)
    return coverage_database_file, coveragercfile


def parse_coverage_file(
    coverage_file_path: Path,
    source_code_path: Path,
) -> tuple[dict[str, dict[str, Any]], CoverageStatus]:
    """Parse a JSON coverage file and return function-level data."""
    with coverage_file_path.open(encoding="utf-8") as f:
        coverage_data = json.load(f)

    candidates = generate_candidates(source_code_path)

    log.debug("Looking for coverage data in %s", " -> ".join(candidates))
    for candidate in candidates:
        try:
            cov: dict[str, dict[str, Any]] = coverage_data["files"][candidate][
                "functions"
            ]
            log.debug(
                "Coverage data found for %s in %s",
                source_code_path,
                candidate,
            )
            status = CoverageStatus.PARSED_SUCCESSFULLY
            break
        except KeyError:
            continue
    else:
        log.debug(
            "No coverage data found for %s in %s",
            source_code_path,
            candidates,
        )
        cov = {}
        status = CoverageStatus.NOT_FOUND
    return cov, status


def fetch_function_coverages(
    function_name: str,
    code_context: CodeOptimizationContext,
    coverage_data: dict[str, dict[str, Any]],
    original_cov_data: dict[str, dict[str, Any]],
) -> tuple[FunctionCoverage, FunctionCoverage | None]:
    """Fetch coverage for the main function and its dependent (if any)."""
    resolved_name = build_fully_qualified_name(function_name, code_context)
    try:
        main_function_coverage = FunctionCoverage(
            name=resolved_name,
            coverage=coverage_data[resolved_name]["summary"][
                "percent_covered"
            ],
            executed_lines=coverage_data[resolved_name]["executed_lines"],
            unexecuted_lines=coverage_data[resolved_name]["missing_lines"],
            executed_branches=coverage_data[resolved_name][
                "executed_branches"
            ],
            unexecuted_branches=coverage_data[resolved_name][
                "missing_branches"
            ],
        )
    except KeyError:
        main_function_coverage = FunctionCoverage(
            name=resolved_name,
            coverage=0,
            executed_lines=[],
            unexecuted_lines=[],
            executed_branches=[],
            unexecuted_branches=[],
        )

    dependent_function = extract_dependent_function(
        function_name, code_context
    )
    dependent_func_coverage = (
        grab_dependent_function_from_coverage_data(
            dependent_function, coverage_data, original_cov_data
        )
        if dependent_function
        else None
    )

    return main_function_coverage, dependent_func_coverage


def aggregate_coverage(
    main_func_coverage: FunctionCoverage,
    dependent_func_coverage: FunctionCoverage | None,
) -> tuple[set[int], set[int]]:
    """Aggregate executed and unexecuted lines across main and dependent functions."""
    total_executed_lines = set(main_func_coverage.executed_lines)
    total_unexecuted_lines = set(main_func_coverage.unexecuted_lines)

    if dependent_func_coverage:
        total_executed_lines.update(dependent_func_coverage.executed_lines)
        total_unexecuted_lines.update(dependent_func_coverage.unexecuted_lines)

    return total_executed_lines, total_unexecuted_lines


def build_coverage_graph(
    main_func_coverage: FunctionCoverage,
    dependent_func_coverage: FunctionCoverage | None,
) -> dict[str, dict[str, Collection[object]]]:
    """Build a per-function graph of executed and unexecuted lines/branches."""
    graph: dict[str, dict[str, Collection[object]]] = {
        main_func_coverage.name: {
            "executed_lines": set(main_func_coverage.executed_lines),
            "unexecuted_lines": set(main_func_coverage.unexecuted_lines),
            "executed_branches": main_func_coverage.executed_branches,
            "unexecuted_branches": main_func_coverage.unexecuted_branches,
        }
    }

    if dependent_func_coverage:
        graph[dependent_func_coverage.name] = {
            "executed_lines": set(dependent_func_coverage.executed_lines),
            "unexecuted_lines": set(dependent_func_coverage.unexecuted_lines),
            "executed_branches": (dependent_func_coverage.executed_branches),
            "unexecuted_branches": (
                dependent_func_coverage.unexecuted_branches
            ),
        }

    return graph


def grab_dependent_function_from_coverage_data(
    dependent_function_name: str,
    coverage_data: dict[str, dict[str, Any]],
    original_cov_data: dict[str, dict[str, Any]],
) -> FunctionCoverage:
    """Grab a dependent function's coverage from the coverage data."""
    try:
        return FunctionCoverage(
            name=dependent_function_name,
            coverage=coverage_data[dependent_function_name]["summary"][
                "percent_covered"
            ],
            executed_lines=coverage_data[dependent_function_name][
                "executed_lines"
            ],
            unexecuted_lines=coverage_data[dependent_function_name][
                "missing_lines"
            ],
            executed_branches=coverage_data[dependent_function_name][
                "executed_branches"
            ],
            unexecuted_branches=coverage_data[dependent_function_name][
                "missing_branches"
            ],
        )
    except KeyError:
        msg = (
            f"Coverage data not found for dependent function"
            f" {dependent_function_name} in the coverage data"
        )
        try:
            files = original_cov_data["files"]
            for file_key in files:
                functions = files[file_key]["functions"]
                for func_key in functions:
                    if func_key == dependent_function_name or (
                        "." in dependent_function_name
                        and func_key.endswith(f".{dependent_function_name}")
                    ):
                        return FunctionCoverage(
                            name=dependent_function_name,
                            coverage=functions[func_key]["summary"][
                                "percent_covered"
                            ],
                            executed_lines=functions[func_key][
                                "executed_lines"
                            ],
                            unexecuted_lines=functions[func_key][
                                "missing_lines"
                            ],
                            executed_branches=functions[func_key][
                                "executed_branches"
                            ],
                            unexecuted_branches=functions[func_key][
                                "missing_branches"
                            ],
                        )
            msg = (
                f"Coverage data not found for dependent function"
                f" {dependent_function_name} in the original"
                f" coverage data"
            )
        except KeyError:
            raise ValueError(msg) from None

    return FunctionCoverage(
        name=dependent_function_name,
        coverage=0,
        executed_lines=[],
        unexecuted_lines=[],
        executed_branches=[],
        unexecuted_branches=[],
    )


def load_coverage_from_sqlite(
    database_path: Path,
    config_path: Path,
    function_name: str,
    code_context: CodeOptimizationContext,
    source_code_path: Path,
) -> CoverageData:
    """Load coverage data from a coverage.py SQLite database."""
    from coverage import (  # noqa: PLC0415
        Coverage,
    )
    from coverage.exceptions import (  # noqa: PLC0415
        NoDataError,
    )
    from coverage.jsonreport import (  # noqa: PLC0415
        JsonReporter,
    )

    cov = Coverage(
        data_file=database_path,
        config_file=config_path,
        data_suffix=True,
        auto_data=True,
        branch=True,
    )

    if not database_path.exists() or not database_path.stat().st_size:
        log.debug(
            "Coverage database %s is empty or does not exist",
            database_path,
        )
        return create_empty_coverage_data(
            source_code_path, function_name, code_context
        )
    cov.load()

    reporter = JsonReporter(cov)
    temp_json_file = database_path.with_suffix(".report.json")
    with temp_json_file.open("w", encoding="utf-8") as f:
        try:
            reporter.report(morfs=[source_code_path.as_posix()], outfile=f)
        except NoDataError:
            log.debug(
                "No coverage data found for %s in %s",
                function_name,
                source_code_path,
            )
            return create_empty_coverage_data(
                source_code_path, function_name, code_context
            )
    with temp_json_file.open() as f:
        original_coverage_data = json.load(f)

    cov_data, status = parse_coverage_file(temp_json_file, source_code_path)

    main_func_coverage, dependent_func_coverage = fetch_function_coverages(
        function_name,
        code_context,
        cov_data,
        original_cov_data=original_coverage_data,
    )

    total_executed_lines, total_unexecuted_lines = aggregate_coverage(
        main_func_coverage, dependent_func_coverage
    )

    total_lines = total_executed_lines | total_unexecuted_lines
    coverage_pct = (
        len(total_executed_lines) / len(total_lines) * 100
        if total_lines
        else 0.0
    )

    functions_being_tested = [main_func_coverage.name]
    if dependent_func_coverage:
        functions_being_tested.append(dependent_func_coverage.name)

    graph = build_coverage_graph(main_func_coverage, dependent_func_coverage)
    temp_json_file.unlink()

    return CoverageData(
        file_path=source_code_path,
        coverage=coverage_pct,
        function_name=function_name,
        functions_being_tested=functions_being_tested,
        graph=graph,
        code_context=code_context,
        main_func_coverage=main_func_coverage,
        dependent_func_coverage=dependent_func_coverage,
        status=status,
    )
