"""Pytest plugin for Codeflash benchmark tracing.

Provides a ``benchmark`` fixture that records end-to-end benchmark
timings alongside per-function trace data collected by
:data:`codeflash_trace`.
"""

from __future__ import annotations

import importlib.util
import os
import sqlite3
import sys
import time
from pathlib import Path
from typing import TYPE_CHECKING, Any

import pytest

from ..test_discovery.linking import module_name_from_file_path
from ._benchmark_tracing import codeflash_trace

if TYPE_CHECKING:
    from .models import BenchmarkKey

PYTEST_BENCHMARK_INSTALLED = (
    importlib.util.find_spec("pytest_benchmark") is not None
)


def pytest_addoption(parser: pytest.Parser) -> None:
    """Register the ``--codeflash-trace`` command-line flag."""
    parser.addoption(
        "--codeflash-trace",
        action="store_true",
        default=False,
        help="Enable Codeflash benchmark tracing",
    )


@pytest.fixture
def benchmark(request: pytest.FixtureRequest) -> Any:
    """Provide the Codeflash benchmark fixture.

    Delegates to :class:`CodeFlashBenchmarkPlugin.Benchmark`
    when tracing is active, otherwise falls back to a no-op
    callable.
    """
    if request.config.getoption(
        "--codeflash-trace",
        False,  # noqa: FBT003
    ):
        return CodeFlashBenchmarkPlugin.Benchmark(request)
    # Passthrough: just call the function directly
    return lambda fn, *a, **kw: fn(*a, **kw)


class CodeFlashBenchmarkPlugin:
    """Pytest plugin that captures benchmark timing data."""

    def __init__(self) -> None:
        """Initialize the plugin with empty state."""
        self._trace_path: str | None = None
        self._connection: sqlite3.Connection | None = None
        self.project_root: str | None = None
        self.benchmark_timings: list[tuple[str, str, int, int]] = []

    def setup(
        self,
        trace_path: str | Path,
        project_root: str | Path,
    ) -> None:
        """Create the benchmark_timings table in the trace database."""
        try:
            self.project_root = str(project_root)
            self._trace_path = str(trace_path)
            self._connection = sqlite3.connect(self._trace_path)
            cur = self._connection.cursor()
            cur.execute("PRAGMA synchronous = OFF")
            cur.execute("PRAGMA journal_mode = MEMORY")
            cur.execute(
                "CREATE TABLE IF NOT EXISTS"
                " benchmark_timings("
                "benchmark_module_path TEXT,"
                " benchmark_function_name TEXT,"
                " benchmark_line_number INTEGER,"
                "benchmark_time_ns INTEGER)"
            )
            self._connection.commit()
            self.close()
        except Exception as e:
            print(f"Database setup error: {e}")
            if self._connection:
                self._connection.close()
                self._connection = None
            raise

    def write_benchmark_timings(self) -> None:
        """Flush buffered benchmark timing rows to the database."""
        if not self.benchmark_timings:
            return
        if self._connection is None:
            self._connection = sqlite3.connect(
                self._trace_path  # type: ignore[arg-type]
            )
        try:
            cur = self._connection.cursor()
            cur.executemany(
                "INSERT INTO benchmark_timings"
                " (benchmark_module_path,"
                " benchmark_function_name,"
                " benchmark_line_number,"
                " benchmark_time_ns)"
                " VALUES (?, ?, ?, ?)",
                self.benchmark_timings,
            )
            self._connection.commit()
            self.benchmark_timings = []
        except Exception as e:
            print(f"Error writing to benchmark timings database: {e}")
            self._connection.rollback()
            raise

    def close(self) -> None:
        """Close the SQLite connection."""
        if self._connection:
            self._connection.close()
            self._connection = None

    @staticmethod
    def get_function_benchmark_timings(
        trace_path: Path,
    ) -> dict[str, dict[BenchmarkKey, int]]:
        """Extract per-function timing data from trace files."""
        from .models import BenchmarkKey  # noqa: PLC0415

        result: dict[str, dict[BenchmarkKey, int]] = {}
        connection = sqlite3.connect(trace_path)
        cursor = connection.cursor()
        try:
            cursor.execute(
                "SELECT module_name, class_name,"
                " function_name,"
                " benchmark_module_path,"
                " benchmark_function_name,"
                " benchmark_line_number,"
                " function_time_ns"
                " FROM benchmark_function_timings"
            )
            for row in cursor.fetchall():
                (
                    module_name,
                    class_name,
                    function_name,
                    benchmark_file,
                    benchmark_func,
                    _benchmark_line,
                    time_ns,
                ) = row
                if class_name:
                    qualified_name = (
                        f"{module_name}.{class_name}.{function_name}"
                    )
                else:
                    qualified_name = f"{module_name}.{function_name}"
                benchmark_key = BenchmarkKey(
                    module_path=benchmark_file,
                    function_name=benchmark_func,
                )
                if qualified_name not in result:
                    result[qualified_name] = {}
                if benchmark_key in result[qualified_name]:
                    result[qualified_name][benchmark_key] += time_ns
                else:
                    result[qualified_name][benchmark_key] = time_ns
        finally:
            connection.close()
        return result

    @staticmethod
    def get_benchmark_timings(
        trace_path: Path,
    ) -> dict[BenchmarkKey, int]:
        """Extract total benchmark timings from trace files."""
        from .models import BenchmarkKey  # noqa: PLC0415

        result: dict[BenchmarkKey, int] = {}
        overhead_by_benchmark: dict[BenchmarkKey, int] = {}
        connection = sqlite3.connect(trace_path)
        cursor = connection.cursor()
        try:
            cursor.execute(
                "SELECT benchmark_module_path,"
                " benchmark_function_name,"
                " benchmark_line_number,"
                " SUM(overhead_time_ns)"
                " FROM benchmark_function_timings"
                " GROUP BY benchmark_module_path,"
                " benchmark_function_name,"
                " benchmark_line_number"
            )
            for row in cursor.fetchall():
                (
                    benchmark_file,
                    benchmark_func,
                    _benchmark_line,
                    total_overhead_ns,
                ) = row
                benchmark_key = BenchmarkKey(
                    module_path=benchmark_file,
                    function_name=benchmark_func,
                )
                overhead_by_benchmark[benchmark_key] = total_overhead_ns or 0

            cursor.execute(
                "SELECT benchmark_module_path,"
                " benchmark_function_name,"
                " benchmark_line_number,"
                " benchmark_time_ns"
                " FROM benchmark_timings"
            )
            for row in cursor.fetchall():
                (
                    benchmark_file,
                    benchmark_func,
                    _benchmark_line,
                    time_ns,
                ) = row
                benchmark_key = BenchmarkKey(
                    module_path=benchmark_file,
                    function_name=benchmark_func,
                )
                overhead = overhead_by_benchmark.get(benchmark_key, 0)
                result[benchmark_key] = time_ns - overhead
        finally:
            connection.close()
        return result

    @pytest.hookimpl
    def pytest_sessionfinish(
        self,
        session: pytest.Session,
        exitstatus: int,
    ) -> None:
        """Flush remaining data and close connections at session end."""
        codeflash_trace.close()
        if self.benchmark_timings:
            self.write_benchmark_timings()
        self.close()

    @staticmethod
    def pytest_collection_modifyitems(
        config: pytest.Config,
        items: list[pytest.Item],
    ) -> None:
        """Skip non-benchmark tests when ``--codeflash-trace`` is active."""
        if not config.getoption("--codeflash-trace"):
            return
        skip_no_benchmark = pytest.mark.skip(
            reason="Test requires benchmark fixture"
        )
        for item in items:
            has_fixture = (
                hasattr(item, "fixturenames")
                and "benchmark" in item.fixturenames
            )
            has_marker = False
            if hasattr(item, "get_closest_marker"):
                marker = item.get_closest_marker("benchmark")
                if marker is not None:
                    has_marker = True
            if not (has_fixture or has_marker):
                item.add_marker(skip_no_benchmark)

    class Benchmark:
        """Callable that executes and times a benchmark function."""

        def __init__(self, request: pytest.FixtureRequest) -> None:
            """Store the pytest request for path resolution."""
            self.request = request

        def __call__(self, func: Any, *args: Any, **kwargs: Any) -> Any:
            """Run *func* as a benchmark, with or without arguments."""
            if args or kwargs:
                return self._run_benchmark(func, *args, **kwargs)

            def wrapped_func(*a: Any, **kw: Any) -> Any:
                """Transparent wrapper returned when no args are given."""
                return func(*a, **kw)

            self._run_benchmark(func)
            return wrapped_func

        def _run_benchmark(self, func: Any, *args: Any, **kwargs: Any) -> Any:
            """Execute and record a single benchmark invocation."""
            node_path = getattr(self.request.node, "path", None) or getattr(
                self.request.node, "fspath", None
            )
            if node_path is None:
                msg = "Unable to determine test file path from pytest node"
                raise RuntimeError(msg)
            try:
                benchmark_module_path = module_name_from_file_path(
                    Path(str(node_path)),
                    Path(str(codeflash_benchmark_plugin.project_root)),
                )
            except ValueError:
                benchmark_module_path = Path(str(node_path)).stem
            benchmark_function_name = self.request.node.name
            line_number = int(str(sys._getframe(2).f_lineno))
            os.environ["CODEFLASH_BENCHMARK_FUNCTION_NAME"] = (
                benchmark_function_name
            )
            os.environ["CODEFLASH_BENCHMARK_MODULE_PATH"] = (
                benchmark_module_path
            )
            os.environ["CODEFLASH_BENCHMARK_LINE_NUMBER"] = str(line_number)
            os.environ["CODEFLASH_BENCHMARKING"] = "True"

            start = time.perf_counter_ns()
            result = func(*args, **kwargs)
            end = time.perf_counter_ns()

            os.environ["CODEFLASH_BENCHMARKING"] = "False"
            codeflash_trace.write_function_timings()
            codeflash_trace.function_call_count = 0
            codeflash_benchmark_plugin.benchmark_timings.append(
                (
                    benchmark_module_path,
                    benchmark_function_name,
                    line_number,
                    end - start,
                )
            )
            return result


codeflash_benchmark_plugin = CodeFlashBenchmarkPlugin()
