"""Trace decorator for benchmark function profiling.

Wraps target functions to record execution times and serialized
arguments into a SQLite database during benchmark runs.
"""

from __future__ import annotations

import functools
import os
import pickle
import sqlite3
import threading
import time
from pathlib import Path
from typing import Any

from ..runtime._picklepatch.pickle_patcher import PicklePatcher

BENCHMARK_TIMINGS_SCHEMA: str = (
    "CREATE TABLE IF NOT EXISTS"
    " benchmark_function_timings("
    "function_name TEXT, class_name TEXT,"
    " module_name TEXT, file_path TEXT,"
    "benchmark_function_name TEXT,"
    " benchmark_module_path TEXT,"
    " benchmark_line_number INTEGER,"
    "function_time_ns INTEGER,"
    " overhead_time_ns INTEGER,"
    " args BLOB, kwargs BLOB)"
)


class CodeflashTrace:
    """Decorator class that traces and profiles function execution."""

    def __init__(self) -> None:
        """Initialize the trace decorator with empty state."""
        self.function_calls_data: list[tuple[Any, ...]] = []
        self.function_call_count: int = 0
        self.pickle_count_limit: int = 1000
        self._connection: sqlite3.Connection | None = None
        self._trace_path: str | None = None
        self._thread_local = threading.local()
        self._thread_local.active_functions = set()  # set[tuple[str, str]]

    def setup(self, trace_path: str) -> None:
        """Create the SQLite database and benchmark_function_timings table."""
        try:
            self._trace_path = 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(BENCHMARK_TIMINGS_SCHEMA)
            self._connection.commit()
        except Exception as e:
            print(f"Database setup error: {e}")
            if self._connection:
                self._connection.close()
                self._connection = None
            raise

    def write_function_timings(self) -> None:
        """Flush buffered function timing rows to the SQLite database."""
        if not self.function_calls_data:
            return
        if self._connection is None and self._trace_path is not None:
            self._connection = sqlite3.connect(self._trace_path)
        try:
            cur = self._connection.cursor()  # type: ignore[union-attr]
            cur.executemany(
                "INSERT INTO benchmark_function_timings"
                "(function_name, class_name, module_name,"
                " file_path,"
                " benchmark_function_name,"
                " benchmark_module_path,"
                " benchmark_line_number,"
                " function_time_ns, overhead_time_ns,"
                " args, kwargs) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                self.function_calls_data,
            )
            self._connection.commit()  # type: ignore[union-attr]
            self.function_calls_data = []
        except Exception as e:
            print(f"Error writing to function timings database: {e}")
            if self._connection:
                self._connection.rollback()
            raise

    def open(self) -> None:
        """Re-open the SQLite connection if it was closed."""
        if self._connection is None:
            self._connection = sqlite3.connect(
                self._trace_path  # type: ignore[arg-type]
            )

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

    def __call__(self, func: Any) -> Any:
        """Decorate *func* to trace its execution during benchmarks."""
        func_id = (func.__module__, func.__name__)

        @functools.wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            if not hasattr(self._thread_local, "active_functions"):
                self._thread_local.active_functions = set()

            if func_id in self._thread_local.active_functions:
                return func(*args, **kwargs)

            self._thread_local.active_functions.add(func_id)
            start_time = time.thread_time_ns()
            result = func(*args, **kwargs)
            end_time = time.thread_time_ns()
            execution_time = end_time - start_time
            self.function_call_count += 1

            if os.environ.get("CODEFLASH_BENCHMARKING", "False") == "False":
                self._thread_local.active_functions.remove(func_id)
                return result

            benchmark_function_name = os.environ.get(
                "CODEFLASH_BENCHMARK_FUNCTION_NAME", ""
            )
            benchmark_module_path = os.environ.get(
                "CODEFLASH_BENCHMARK_MODULE_PATH", ""
            )
            benchmark_line_number = os.environ.get(
                "CODEFLASH_BENCHMARK_LINE_NUMBER", ""
            )
            class_name = ""
            qualname = func.__qualname__
            if "." in qualname:
                class_name = qualname.split(".")[0]

            normalized_file_path = Path(func.__code__.co_filename).as_posix()

            if self.function_call_count > self.pickle_count_limit:
                self._thread_local.active_functions.remove(func_id)
                overhead_time = time.thread_time_ns() - end_time
                self.function_calls_data.append(
                    (
                        func.__name__,
                        class_name,
                        func.__module__,
                        normalized_file_path,
                        benchmark_function_name,
                        benchmark_module_path,
                        benchmark_line_number,
                        execution_time,
                        overhead_time,
                        None,
                        None,
                    )
                )
                return result

            try:
                pickled_args = PicklePatcher.dumps(
                    args,
                    protocol=pickle.HIGHEST_PROTOCOL,
                )
                pickled_kwargs = PicklePatcher.dumps(
                    kwargs,
                    protocol=pickle.HIGHEST_PROTOCOL,
                )
            except Exception:
                self._thread_local.active_functions.remove(func_id)
                overhead_time = time.thread_time_ns() - end_time
                self.function_calls_data.append(
                    (
                        func.__name__,
                        class_name,
                        func.__module__,
                        normalized_file_path,
                        benchmark_function_name,
                        benchmark_module_path,
                        benchmark_line_number,
                        execution_time,
                        overhead_time,
                        None,
                        None,
                    )
                )
                return result

            if len(self.function_calls_data) > 100:
                self.write_function_timings()

            self._thread_local.active_functions.remove(func_id)
            overhead_time = time.thread_time_ns() - end_time
            self.function_calls_data.append(
                (
                    func.__name__,
                    class_name,
                    func.__module__,
                    normalized_file_path,
                    benchmark_function_name,
                    benchmark_module_path,
                    benchmark_line_number,
                    execution_time,
                    overhead_time,
                    pickled_args,
                    pickled_kwargs,
                )
            )
            return result

        return wrapper


codeflash_trace = CodeflashTrace()
