"""Rank and filter functions by profiling-derived addressable time."""

from __future__ import annotations

import logging
import sqlite3
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionToOptimize

log = logging.getLogger(__name__)

DEFAULT_IMPORTANCE_THRESHOLD: float = 0.001

PYTEST_FILE_PATTERNS: frozenset[str] = frozenset(
    {
        "<frozen",
        "<string>",
        "_pytest/",
        "pytest",
        "pluggy/",
        "_pydev",
        "runpy.py",
    }
)

PYTEST_FUNC_PATTERNS: frozenset[str] = frozenset(
    {
        "pytest_",
        "_pytest",
        "runtest",
    }
)


def is_pytest_infrastructure(filename: str, function_name: str) -> bool:
    """Check if a function is part of pytest infrastructure."""
    for pattern in PYTEST_FILE_PATTERNS:
        if pattern in filename:
            return True
    return any(
        pattern in function_name.lower() for pattern in PYTEST_FUNC_PATTERNS
    )


class FunctionRanker:
    """Rank and filter functions by addressable time.

    Addressable time is calculated as::

        own_time + (time_in_callees / call_count)

    This prioritises functions that are computationally
    heavy or that make expensive calls to other functions.

    Functions below an importance threshold (fraction of
    total file runtime) are filtered out before ranking.
    """

    def __init__(self, trace_file_path: Path) -> None:
        """Initialize from a SQLite trace file and build the function stats index."""
        self.trace_file_path = trace_file_path
        self._function_stats: dict[str, dict[str, Any]] = {}
        self._function_stats_by_name: dict[
            str, list[tuple[str, dict[str, Any]]]
        ] = {}
        self._load_function_stats()

        # Build index: function_name -> [(key, stats)]
        for key, stats in self._function_stats.items():
            func_name = stats.get("function_name")
            if func_name:
                self._function_stats_by_name.setdefault(func_name, []).append(
                    (key, stats)
                )

    def _load_function_stats(self) -> None:
        """Load and process function stats from the SQLite trace file."""
        try:
            con = sqlite3.connect(self.trace_file_path.as_posix())
            cur = con.cursor()
            pdata = cur.execute("SELECT * FROM pstats").fetchall()
            con.close()

            pytest_filtered_count = 0
            for (
                filename,
                line_number,
                function,
                class_name,
                call_count,
                _num_callers,
                total_time_ns,
                cumulative_time_ns,
                _callers,
            ) in pdata:
                if call_count <= 0:
                    continue

                # Build qualified name from class_name + function columns
                qualified_name = (
                    f"{class_name}.{function}" if class_name else function
                )

                if is_pytest_infrastructure(filename, qualified_name):
                    pytest_filtered_count += 1
                    continue

                # Parse function name to handle methods within classes
                base_function_name = function

                # Calculate own time (total time - time spent in subcalls)
                own_time_ns = total_time_ns
                time_in_callees_ns = cumulative_time_ns - total_time_ns

                # Addressable = own + avg callee time
                addressable_time_ns = own_time_ns + (
                    time_in_callees_ns / call_count
                )

                function_key = f"{filename}:{qualified_name}"
                self._function_stats[function_key] = {
                    "filename": filename,
                    "function_name": base_function_name,
                    "qualified_name": qualified_name,
                    "class_name": class_name,
                    "line_number": line_number,
                    "call_count": call_count,
                    "own_time_ns": own_time_ns,
                    "cumulative_time_ns": cumulative_time_ns,
                    "time_in_callees_ns": time_in_callees_ns,
                    "addressable_time_ns": addressable_time_ns,
                }

            log.debug(
                "Loaded timing stats for %d functions from trace "
                "(filtered %d pytest infrastructure functions)",
                len(self._function_stats),
                pytest_filtered_count,
            )

        except (OSError, sqlite3.Error):
            log.warning(
                "Failed to process function stats from trace file %s",
                self.trace_file_path,
                exc_info=True,
            )
            self._function_stats = {}

    def get_function_stats_summary(
        self, function_to_optimize: FunctionToOptimize
    ) -> dict[str, Any] | None:
        """Look up profiling stats for a function.

        Returns the stats dict if found, or *None* if the function was not
        recorded in the trace.
        """
        target_filename = function_to_optimize.file_path.name
        candidates = self._function_stats_by_name.get(
            function_to_optimize.function_name
        )
        if not candidates:
            log.debug(
                "Could not find stats for function %s in file %s",
                function_to_optimize.function_name,
                target_filename,
            )
            return None

        for key, stats in candidates:
            if key.endswith(f"/{target_filename}") or target_filename in key:
                return stats

        log.debug(
            "Could not find stats for function %s in file %s",
            function_to_optimize.function_name,
            target_filename,
        )
        return None

    def get_function_addressable_time(
        self, function_to_optimize: FunctionToOptimize
    ) -> float:
        """Get the addressable time in nanoseconds for a function."""
        stats = self.get_function_stats_summary(function_to_optimize)
        return stats["addressable_time_ns"] if stats else 0.0

    def rank_functions(
        self,
        functions_to_optimize: list[FunctionToOptimize],
    ) -> list[FunctionToOptimize]:
        """Rank and filter functions by addressable time."""
        if not self._function_stats:
            log.warning("No function stats available to rank functions.")
            return []

        # Total time from same file(s) as target functions
        if functions_to_optimize:
            target_files = {
                func.file_path.name for func in functions_to_optimize
            }
            total_program_time = sum(
                s["own_time_ns"]
                for s in self._function_stats.values()
                if s.get("own_time_ns", 0) > 0
                and any(
                    str(s.get("filename", "")).endswith("/" + target_file)
                    or s.get("filename") == target_file
                    for target_file in target_files
                )
            )
            log.debug(
                "Using file-relative importance for %d file(s): %s. "
                "Total file time: %d ns",
                len(target_files),
                target_files,
                total_program_time,
            )
        else:
            total_program_time = sum(
                s["own_time_ns"]
                for s in self._function_stats.values()
                if s.get("own_time_ns", 0) > 0
            )

        if total_program_time == 0:
            log.warning(
                "Total program time is zero; "
                "cannot determine function importance."
            )
            functions_to_rank = functions_to_optimize
        else:
            functions_to_rank = []
            for func in functions_to_optimize:
                func_stats = self.get_function_stats_summary(func)
                if func_stats and func_stats.get("addressable_time_ns", 0) > 0:
                    importance = (
                        func_stats["addressable_time_ns"] / total_program_time
                    )
                    if importance >= DEFAULT_IMPORTANCE_THRESHOLD:
                        functions_to_rank.append(func)
                    else:
                        log.debug(
                            "Filtering out function %s with importance "
                            "%.2f%% (below threshold %.2f%%)",
                            func.qualified_name,
                            importance * 100,
                            DEFAULT_IMPORTANCE_THRESHOLD * 100,
                        )

            log.info(
                "Filtered down to %d important functions "
                "from %d total functions",
                len(functions_to_rank),
                len(functions_to_optimize),
            )

        ranked = sorted(
            functions_to_rank,
            key=self.get_function_addressable_time,
            reverse=True,
        )
        log.debug(
            "Function ranking order: %s",
            [
                f"{func.function_name} "
                f"(addressable_time={self.get_function_addressable_time(func):.2f}ns)"
                for func in ranked
            ],
        )
        return ranked
