"""Capture decorator for recording __init__ instance state during tests."""

# ruff: noqa: T201, BLE001, C901, FBT001, FBT002
from __future__ import annotations

# This file should not have any dependencies on codeflash
import functools
import gc
import inspect
import os
import sqlite3
import time
import warnings
from pathlib import Path
from typing import Any, Callable

import dill
from dill import PicklingWarning

from codeflash_python._model import VerificationType
from codeflash_python.runtime._picklepatch.pickle_patcher import PicklePatcher

warnings.filterwarnings("ignore", category=PicklingWarning)


def get_test_info_from_stack(
    tests_root: str,
) -> tuple[str, str | None, str | None, str]:
    """Extract test information by walking the call stack from the current frame."""
    test_module_name = ""
    test_class_name: str | None = None
    test_name: str | None = None
    line_id = ""

    # Get current frame and skip our own function's frame
    frame = inspect.currentframe()
    if frame is not None:
        frame = frame.f_back

    # Walk the stack
    while frame is not None:
        function_name = frame.f_code.co_name
        filename = frame.f_code.co_filename
        lineno = frame.f_lineno

        # Check if function name indicates a test (e.g., starts with "test_")
        if function_name.startswith("test_"):
            test_name = function_name
            test_module = inspect.getmodule(frame)
            if test_module is not None and hasattr(test_module, "__name__"):
                test_module_name = test_module.__name__
            line_id = str(lineno)

            # Check if it's a method in a class
            if (
                "self" in frame.f_locals
                and hasattr(frame.f_locals["self"], "__class__")
                and hasattr(frame.f_locals["self"].__class__, "__name__")
            ):
                test_class_name = frame.f_locals["self"].__class__.__name__
            break

        # Check for instantiation on the module level
        if (
            "__name__" in frame.f_globals
            and frame.f_globals["__name__"].split(".")[-1].startswith("test_")
            and Path(filename).resolve().is_relative_to(Path(tests_root))
            and function_name == "<module>"
        ):
            test_module_name = frame.f_globals["__name__"]
            line_id = str(lineno)

            #     # Check if it's a method in a class
            if (
                "self" in frame.f_locals
                and hasattr(frame.f_locals["self"], "__class__")
                and hasattr(frame.f_locals["self"].__class__, "__name__")
            ):
                test_class_name = frame.f_locals["self"].__class__.__name__
            break

        # Go to the previous frame
        frame = frame.f_back

    # If stack walking didn't find test info, fall back to environment variables
    if not test_name:
        env_test_function = os.environ.get("CODEFLASH_TEST_FUNCTION")
        if env_test_function:
            test_name = env_test_function
            if not test_module_name:
                test_module_name = os.environ.get("CODEFLASH_TEST_MODULE", "")
            if not test_class_name:
                env_class = os.environ.get("CODEFLASH_TEST_CLASS")
                test_class_name = env_class or None

    return test_module_name, test_class_name, test_name, line_id


def codeflash_capture(
    function_name: str,
    tmp_dir_path: str,
    tests_root: str,
    is_fto: bool = False,
) -> Callable[..., Any]:
    """Define a decorator to instrument the init function, collect test info, and capture the instance state."""

    def decorator(wrapped: Callable[..., Any]) -> Callable[..., Any]:
        """Wrap the __init__ function with timing and state capture logic."""

        @functools.wraps(wrapped)
        def wrapper(*args: Any, **kwargs: Any) -> None:
            """Execute wrapped __init__, recording timing and instance state."""
            # Dynamic information retrieved from stack
            test_module_name, test_class_name, test_name, line_id = (
                get_test_info_from_stack(tests_root)
            )

            # Get env variables
            loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
            codeflash_iteration = os.environ["CODEFLASH_TEST_ITERATION"]

            # Create test_id
            test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{loop_index}"

            # Initialize index tracking if needed, handles multiple instances created in the same test line number
            if not hasattr(wrapper, "index"):
                wrapper.index = {}  # type: ignore[attr-defined]

            # Update index for this test
            if test_id in wrapper.index:  # type: ignore[attr-defined]
                wrapper.index[test_id] += 1  # type: ignore[attr-defined]
            else:
                wrapper.index[test_id] = 0  # type: ignore[attr-defined]

            codeflash_test_index = wrapper.index[test_id]  # type: ignore[attr-defined]

            # Generate invocation id
            invocation_id = f"{line_id}_{codeflash_test_index}"
            test_stdout_tag = f"{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}"
            print(f"!$######{test_stdout_tag}######$!")
            # Connect to sqlite
            codeflash_con = sqlite3.connect(
                f"{tmp_dir_path}_{codeflash_iteration}.sqlite"
            )
            codeflash_cur = codeflash_con.cursor()

            # Record timing information
            exception = None
            gc.disable()
            try:
                counter = time.perf_counter_ns()
                wrapped(*args, **kwargs)
                codeflash_duration = time.perf_counter_ns() - counter
            except Exception as e:
                codeflash_duration = time.perf_counter_ns() - counter
                exception = e
            finally:
                gc.enable()
            print(f"!######{test_stdout_tag}######!")
            # Capture instance state after initialization
            # self is always the first argument, this is ensured during instrumentation
            instance = args[0]
            if hasattr(instance, "__dict__"):
                instance_state = instance.__dict__
            elif hasattr(instance, "__slots__"):
                # For classes using __slots__, capture slot values
                instance_state = {
                    slot: getattr(instance, slot, None)
                    for slot in instance.__slots__
                    if hasattr(instance, slot)
                }
            else:
                # For C extension types or other special classes (e.g., Playwright's Page),
                # capture all non-private, non-callable attributes
                instance_state = {
                    attr: getattr(instance, attr)
                    for attr in dir(instance)
                    if not attr.startswith("_")
                    and not callable(getattr(instance, attr, None))
                }
            codeflash_cur.execute(
                "CREATE TABLE IF NOT EXISTS test_results (test_module_path TEXT, test_class_name TEXT, test_function_name TEXT, function_getting_tested TEXT, loop_index INTEGER, iteration_id TEXT, runtime INTEGER, return_value BLOB, verification_type TEXT)"
            )

            # Write to sqlite
            pickled_return_value = (
                dill.dumps(exception)
                if exception
                else PicklePatcher.dumps(instance_state)
            )
            codeflash_cur.execute(
                "INSERT INTO test_results VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    test_module_name,
                    test_class_name,
                    test_name,
                    function_name,
                    loop_index,
                    invocation_id,
                    codeflash_duration,
                    pickled_return_value,
                    VerificationType.INIT_STATE_FTO
                    if is_fto
                    else VerificationType.INIT_STATE_HELPER,
                ),
            )
            codeflash_con.commit()
            if exception:
                raise exception

        return wrapper

    return decorator
