"""Tracing and profiling module for capturing function calls."""

from __future__ import annotations

import contextlib
import datetime
import json
import logging
import os
import pickle
import sqlite3
import sys
import threading
import time
from collections import defaultdict
from pathlib import Path
from typing import TYPE_CHECKING, Any, ClassVar

from ..test_discovery.linking import module_name_from_file_path
from ._file_filtering import filter_files_optimized
from ._trace_db import FUNCTION_CALLS_SCHEMA, TOTAL_TIME_SCHEMA
from ._trace_models import FakeCode, FakeFrame, TracedFunction

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence
    from types import FrameType, TracebackType

log = logging.getLogger(__name__)


class Tracer:
    """Profile and trace Python function calls via sys.setprofile.

    Use as a context manager. Stores call data and profiling statistics
    in a SQLite database at *output_file*.
    """

    used_once: ClassVar[bool] = False

    def __init__(  # noqa: PLR0913
        self,
        *,
        project_root: Path,
        module_root: Path,
        tests_root: Path,
        output_file: Path,
        functions: list[str] | None = None,
        ignore_paths: Sequence[Path] | None = None,
        max_function_count: int = 256,
        timeout: float | None = None,
        command: str = "",
        file_filter: Callable[[Path], bool] | None = None,
    ) -> None:
        """Initialize the tracer with project paths, output file, and profiler state."""
        if functions is None:
            functions = []
        self.disable = False
        self._db_lock: threading.Lock | None = None
        if os.environ.get("CODEFLASH_TRACER_DISABLE", "0") == "1":
            log.warning(
                "Tracer disabled by environment variable"
                " CODEFLASH_TRACER_DISABLE."
            )
            self.disable = True
            return
        if sys.getprofile() is not None or sys.gettrace() is not None:
            log.warning(
                "Another profiler or debugger is active — Tracer is disabled."
            )
            self.disable = True
            return

        self._db_lock = threading.Lock()
        self.con: sqlite3.Connection | None = None
        self.functions = functions
        self.function_modules: list[TracedFunction] = []
        self._function_qualified_names: list[str] = []
        self.function_count: dict[str, int] = defaultdict(int)
        self.current_file_path = str(Path(__file__).resolve())
        self.ignored_qualified_functions: set[str] = {
            f"{self.current_file_path}:Tracer.__exit__",
            f"{self.current_file_path}:Tracer.__enter__",
        }
        self.max_function_count = max_function_count
        self.project_root = project_root
        self.project_root_str = str(project_root) + os.sep
        self.module_root = module_root
        self.tests_root = tests_root
        self.ignore_paths = list(ignore_paths) if ignore_paths else []
        self.file_filter = file_filter
        self.output_file = output_file
        self.ignored_functions = {
            "<listcomp>",
            "<genexpr>",
            "<dictcomp>",
            "<setcomp>",
            "<lambda>",
            "<module>",
        }

        assert timeout is None or timeout > 0  # noqa: S101
        self.timeout = timeout
        self.next_insert = 1000
        self._trace_count = 0
        self.path_cache: dict[str, tuple[str, bool]] = {}

        # Profiler variables (adapted from cProfile)
        self.bias = 0
        self.timings: dict[tuple[Any, ...], tuple[Any, ...]] = {}
        self.cur: tuple[Any, ...] | None = None
        self.start_time: float | None = None
        self.timer = time.process_time_ns
        self.total_tt = 0
        self.stats: dict[tuple[Any, ...], tuple[Any, ...]] = {}
        self.simulate_call("profiler")
        self.t = self.timer()

        self.command = command
        self.c_func_name = ""

    @property
    def trace_count(self) -> int:
        """Return the number of traced calls."""
        return self._trace_count

    @property
    def traced_functions(self) -> list[TracedFunction]:
        """Return a copy of the discovered functions."""
        return list(self.function_modules)

    def __enter__(self) -> None:
        """Set up the trace database and install the profiler."""
        if self.disable:
            return
        if Tracer.used_once:
            log.warning(
                "Tracer already used once in this process — disabling."
            )
            self.disable = True
            return
        Tracer.used_once = True

        Path(self.output_file).unlink(missing_ok=True)

        self.con = sqlite3.connect(self.output_file, check_same_thread=False)
        cur = self.con.cursor()
        cur.execute("PRAGMA synchronous = OFF")
        cur.execute("PRAGMA journal_mode = WAL")
        cur.execute(FUNCTION_CALLS_SCHEMA)
        cur.execute("CREATE TABLE metadata(key TEXT PRIMARY KEY, value TEXT)")
        cur.execute(
            "INSERT INTO metadata VALUES (?, ?)",
            ("command", self.command),
        )
        cur.execute(
            "INSERT INTO metadata VALUES (?, ?)",
            (
                "functions_filter",
                json.dumps(self.functions) if self.functions else None,
            ),
        )
        cur.execute(
            "INSERT INTO metadata VALUES (?, ?)",
            (
                "timestamp",
                datetime.datetime.now(datetime.timezone.utc).isoformat(),
            ),
        )
        cur.execute(
            "INSERT INTO metadata VALUES (?, ?)",
            ("project_root", str(self.project_root)),
        )

        frame = sys._getframe(0)  # noqa: SLF001
        self.dispatch["call"](self, frame, 0)
        self.start_time = time.time()
        sys.setprofile(self.trace_callback)
        threading.setprofile(self.trace_callback)

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Uninstall the profiler, persist stats, and close the database."""
        if self.disable or self._db_lock is None:
            return
        sys.setprofile(None)
        threading.setprofile(None)

        try:
            with self._db_lock:
                if self.con is None:
                    return

                self.con.commit()
                self.create_stats()

                cur = self.con.cursor()
                cur.execute(
                    "CREATE TABLE pstats ("
                    "filename TEXT, line_number INTEGER, "
                    "function TEXT, class_name TEXT, "
                    "call_count_nonrecursive INTEGER, "
                    "num_callers INTEGER, "
                    "total_time_ns INTEGER, "
                    "cumulative_time_ns INTEGER, "
                    "callers BLOB)"
                )
                for func, (
                    cc,
                    nc,
                    tt,
                    ct,
                    callers,
                ) in self.stats.items():
                    remapped_callers = [
                        {"key": k, "value": v} for k, v in callers.items()
                    ]
                    cur.execute(
                        "INSERT INTO pstats "
                        "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                        (
                            str(Path(func[0]).resolve()),
                            func[1],
                            func[2],
                            func[3],
                            cc,
                            nc,
                            tt,
                            ct,
                            json.dumps(remapped_callers),
                        ),
                    )
                self.con.commit()

                self.total_tt = sum(
                    tt for _, _, tt, _, _ in self.stats.values()
                )

                cur.execute(TOTAL_TIME_SCHEMA)
                cur.execute(
                    "INSERT INTO total_time VALUES (?)",
                    (self.total_tt,),
                )
                self.con.commit()
                self.con.close()
                self.con = None

            # Filter out functions where no calls were captured
            paired = zip(
                self.function_modules,
                self._function_qualified_names,
            )
            self.function_modules = [
                func
                for func, qname in paired
                if self.function_count[qname] > 0
            ]
        finally:
            Tracer.used_once = False

    def tracer_logic(  # noqa: C901, PLR0911, PLR0912, PLR0915
        self,
        frame: FrameType,
        event: str,
    ) -> None:
        """Record a call event into the trace database."""
        if event != "call":
            return
        if (
            self.timeout is not None
            and self.start_time is not None
            and (time.time() - self.start_time) > self.timeout
        ):
            sys.setprofile(None)
            threading.setprofile(None)
            log.info(
                "Tracer timeout reached at %s seconds.",
                self.timeout,
            )
            return
        if self.disable or self._db_lock is None or self.con is None:
            return

        code = frame.f_code
        if code.co_name in self.ignored_functions:
            return

        co_filename = code.co_filename
        if co_filename in self.path_cache:
            file_name, is_valid = self.path_cache[co_filename]
            if not is_valid:
                return
        else:
            resolved = os.path.realpath(co_filename)
            is_valid = (
                resolved.startswith(self.project_root_str)
                and Path(resolved).exists()
            )
            self.path_cache[co_filename] = (resolved, is_valid)
            if not is_valid:
                return
            file_name = resolved

        if self.functions and code.co_name not in self.functions:
            return

        class_name = None
        method_type = None
        arguments = frame.f_locals
        try:
            self_arg = arguments.get("self")
            if self_arg is not None:
                try:
                    class_name = self_arg.__class__.__name__
                    method_type = "instance"
                except AttributeError:
                    cls_arg = arguments.get("cls")
                    if cls_arg is not None:
                        with contextlib.suppress(AttributeError):
                            class_name = cls_arg.__name__
                            method_type = "classmethod"
            else:
                cls_arg = arguments.get("cls")
                if cls_arg is not None:
                    with contextlib.suppress(AttributeError):
                        class_name = cls_arg.__name__
                        method_type = "classmethod"
        except Exception:  # noqa: BLE001
            log.warning(
                "Failed to extract arguments for %s", file_name, exc_info=True
            )
            return

        # Static methods: class name from co_qualname
        if class_name is None and "." in getattr(code, "co_qualname", ""):
            qualname_parts = code.co_qualname.split(".")
            if len(qualname_parts) >= 2:  # noqa: PLR2004
                class_name = qualname_parts[-2]
                method_type = "staticmethod"

        try:
            function_qualified_name = f"{file_name}:{code.co_qualname}"
        except AttributeError:
            function_qualified_name = (
                f"{file_name}:"
                f"{(class_name + '.' if class_name else '')}"
                f"{code.co_name}"
            )

        if function_qualified_name in self.ignored_qualified_functions:
            return

        if function_qualified_name not in self.function_count:
            self.function_count[function_qualified_name] = 1
            file_path = Path(file_name)

            # File filtering
            if self.file_filter is not None:
                file_valid = self.file_filter(file_path)
            else:
                file_valid = filter_files_optimized(
                    file_path=file_path,
                    tests_root=self.tests_root,
                    ignore_paths=self.ignore_paths,
                    module_root=self.module_root,
                )
            if not file_valid:
                self.ignored_qualified_functions.add(function_qualified_name)
                return

            # Determine is_top_level from co_qualname
            qualname = getattr(code, "co_qualname", code.co_name)
            is_top_level = "<locals>" not in qualname

            self.function_modules.append(
                TracedFunction(
                    function_name=code.co_name,
                    file_name=file_path,
                    module_name=module_name_from_file_path(
                        file_path, self.project_root
                    ),
                    class_name=class_name,
                    line_no=code.co_firstlineno,
                    method_type=method_type,
                    is_top_level=is_top_level,
                )
            )
            self._function_qualified_names.append(function_qualified_name)
        else:
            self.function_count[function_qualified_name] += 1
            if (
                self.function_count[function_qualified_name]
                >= self.max_function_count
            ):
                self.ignored_qualified_functions.add(function_qualified_name)
                return

        with self._db_lock:
            if self.con is None:
                return

            cur = self.con.cursor()
            t_ns = time.perf_counter_ns()
            original_recursion_limit = sys.getrecursionlimit()
            try:
                sys.setrecursionlimit(10000)
                arguments_copy = dict(arguments.items())
                if (
                    class_name
                    and code.co_name == "__init__"
                    and "self" in arguments_copy
                ):
                    del arguments_copy["self"]
                local_vars = pickle.dumps(
                    arguments_copy,
                    protocol=pickle.HIGHEST_PROTOCOL,
                )
                sys.setrecursionlimit(original_recursion_limit)
            except Exception:  # noqa: BLE001
                log.warning(
                    "Failed to pickle arguments for %s",
                    function_qualified_name,
                    exc_info=True,
                )
                self.function_count[function_qualified_name] -= 1
                sys.setrecursionlimit(original_recursion_limit)
                return

            cur.execute(
                "INSERT INTO function_calls VALUES(?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    event,
                    code.co_name,
                    class_name,
                    file_name,
                    frame.f_lineno,
                    frame.f_back.__hash__(),
                    t_ns,
                    local_vars,
                ),
            )
            self._trace_count += 1
            self.next_insert -= 1
            if self.next_insert == 0:
                self.next_insert = 1000
                self.con.commit()

    def trace_callback(
        self,
        frame: FrameType,
        event: str,
        arg: object,
    ) -> None:
        """Dispatch a profiler event."""
        timer = self.timer
        t = timer() - self.t - self.bias
        if event == "c_call":
            self.c_func_name = getattr(arg, "__name__", "")

        prof_success = bool(self.dispatch[event](self, frame, t))
        self.tracer_logic(frame, event)
        if prof_success:
            self.t = timer()
        else:
            self.t = timer() - t

    def trace_dispatch_call(self, frame: FrameType, t: int) -> int:
        """Handle a 'call' profiler event."""
        try:
            if self.cur and frame.f_back is not self.cur[-2]:
                _rpt, _rit, _ret, _rfn, rframe, _rcur = self.cur
                if (
                    not isinstance(rframe, FakeFrame)
                    and hasattr(rframe, "f_back")
                    and hasattr(frame, "f_back")
                    and rframe.f_back is frame.f_back
                ):
                    self.trace_dispatch_return(rframe, 0)

            fcode = frame.f_code
            frame_locals = frame.f_locals
            class_name = None
            try:
                if (
                    "self" in frame_locals
                    and hasattr(frame_locals["self"], "__class__")
                    and hasattr(
                        frame_locals["self"].__class__,
                        "__name__",
                    )
                ):
                    class_name = frame_locals["self"].__class__.__name__
                elif "cls" in frame_locals and hasattr(
                    frame_locals["cls"], "__name__"
                ):
                    class_name = frame_locals["cls"].__name__
            except Exception:  # noqa: BLE001
                log.warning(
                    "Failed to resolve class name in profiler", exc_info=True
                )

            fn = (
                fcode.co_filename,
                fcode.co_firstlineno,
                fcode.co_name,
                class_name,
            )
            self.cur = (t, 0, 0, fn, frame, self.cur)
            timings = self.timings
            if fn in timings:
                cc, ns, tt, ct, callers = timings[fn]
                timings[fn] = cc, ns + 1, tt, ct, callers
            else:
                timings[fn] = 0, 0, 0, 0, {}
        except Exception:  # noqa: BLE001
            log.warning(
                "Profiler dispatch_call failed for frame", exc_info=True
            )
            return 0
        else:
            return 1

    def trace_dispatch_exception(self, frame: FrameType, t: int) -> int:
        """Handle an 'exception' profiler event."""
        rpt, rit, ret, rfn, rframe, rcur = self.cur  # type: ignore[misc]
        if (rframe is not frame) and rcur:
            return self.trace_dispatch_return(rframe, t)
        self.cur = rpt, rit + t, ret, rfn, rframe, rcur
        return 1

    def trace_dispatch_c_call(self, frame: FrameType, t: int) -> int:
        """Handle a 'c_call' profiler event."""
        fn = ("", 0, self.c_func_name, None)
        self.cur = (t, 0, 0, fn, frame, self.cur)
        timings = self.timings
        if fn in timings:
            cc, ns, tt, ct, callers = timings[fn]
            timings[fn] = cc, ns + 1, tt, ct, callers
        else:
            timings[fn] = 0, 0, 0, 0, {}
        return 1

    def trace_dispatch_return(self, frame: FrameType, t: int) -> int:
        """Handle a 'return' profiler event."""
        if not self.cur or not self.cur[-2]:
            return 0
        if frame is not self.cur[-2]:
            if (
                hasattr(frame, "f_back")
                and hasattr(self.cur[-2], "f_back")
                and frame is self.cur[-2].f_back
            ):
                self.trace_dispatch_return(self.cur[-2], 0)
            else:
                return 0

        rpt, rit, ret, rfn, frame, rcur = self.cur
        if not rcur:
            return 0
        rit = rit + t
        frame_total = rit + ret
        ppt, pit, pet, pfn, pframe, pcur = rcur
        self.cur = (
            ppt,
            pit + rpt,
            pet + frame_total,
            pfn,
            pframe,
            pcur,
        )

        timings = self.timings
        if rfn not in timings:
            timings[rfn] = 0, 0, 0, 0, {}
        cc, ns, tt, ct, callers = timings[rfn]
        if not ns:
            ct = ct + frame_total
            cc = cc + 1
        if pfn in callers:
            callers[pfn] = callers[pfn] + 1
        else:
            callers[pfn] = 1
        timings[rfn] = cc, ns - 1, tt + rit, ct, callers
        return 1

    dispatch: ClassVar[dict[str, Callable[..., int]]] = {
        "call": trace_dispatch_call,
        "exception": trace_dispatch_exception,
        "return": trace_dispatch_return,
        "c_call": trace_dispatch_c_call,
        "c_exception": trace_dispatch_return,
        "c_return": trace_dispatch_return,
    }

    def simulate_call(self, name: str) -> None:
        """Push a synthetic call frame onto the profiler stack."""
        code = FakeCode("profiler", 0, name)
        pframe = self.cur[-2] if self.cur else None
        frame = FakeFrame(code, pframe)
        self.dispatch["call"](self, frame, 0)

    def simulate_cmd_complete(self) -> None:
        """Drain all pending return events from the profiler."""
        get_time = self.timer
        t = get_time() - self.t
        while self.cur and self.cur[-1]:
            self.dispatch["return"](self, self.cur[-2], t)
            t = 0
        self.t = get_time() - t

    def create_stats(self) -> None:
        """Finish profiling and compute statistics."""
        self.simulate_cmd_complete()
        self.snapshot_stats()

    def snapshot_stats(self) -> None:
        """Snapshot the current timings into *self.stats*."""
        self.stats = {}
        for func, (
            cc,
            _ns,
            tt,
            ct,
            caller_dict,
        ) in list(self.timings.items()):
            callers = caller_dict.copy()
            nc = 0
            for callcnt in callers.values():
                nc += callcnt
            self.stats[func] = cc, nc, tt, ct, callers

    def runctx(
        self,
        cmd: str,
        global_vars: dict[str, Any],
        local_vars: dict[str, Any] | None,
    ) -> Tracer:
        """Run *cmd* under the tracer and return self."""
        self.__enter__()
        try:
            exec(cmd, global_vars, local_vars)  # noqa: S102
        finally:
            self.__exit__(None, None, None)
        return self
