"""
Self-contained async instrumentation decorators for codeflash testing.

This module is copied to test project directories at runtime, so it
must have zero imports from the codeflash_python package.
"""

# ruff: noqa: BLE001
from __future__ import annotations

import asyncio
import atexit
import contextvars
import gc
import io
import os
import sqlite3
import sys
import time
from enum import Enum
from functools import wraps
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Any, Callable, TypeVar

import dill as pickle


class VerificationType(str, Enum):
    """Type of correctness verification for captured test data."""

    FUNCTION_CALL = "function_call"
    INIT_STATE_FTO = "init_state_fto"
    INIT_STATE_HELPER = "init_state_helper"


F = TypeVar("F", bound=Callable[..., Any])


def get_run_tmp_file(file_path: Path) -> Path:
    """Return a path inside a persistent per-run temporary directory.

    Uses ``CODEFLASH_RUN_TMPDIR`` if set (subprocess case), otherwise
    creates a new tmpdir and exports the env var so child processes
    share the same directory.
    """
    if not hasattr(get_run_tmp_file, "tmpdir"):
        env_dir = os.environ.get("CODEFLASH_RUN_TMPDIR")
        if env_dir and Path(env_dir).is_dir():
            get_run_tmp_file.tmpdir = env_dir  # type: ignore[attr-defined]
        else:
            td = TemporaryDirectory(prefix="codeflash_")
            get_run_tmp_file.tmpdir = td.name  # type: ignore[attr-defined]
            get_run_tmp_file.td_ref = td  # type: ignore[attr-defined]
            os.environ["CODEFLASH_RUN_TMPDIR"] = td.name
    return Path(get_run_tmp_file.tmpdir) / file_path  # type: ignore[attr-defined]


def extract_test_context_from_env() -> tuple[str, str | None, str]:
    """
    Read test module, class, and function names from env vars.

    Raises RuntimeError when required variables are missing.
    """
    test_module = os.environ["CODEFLASH_TEST_MODULE"]
    test_class = os.environ.get("CODEFLASH_TEST_CLASS", None)
    test_function = os.environ["CODEFLASH_TEST_FUNCTION"]
    if test_module and test_function:
        return (
            test_module,
            test_class or None,
            test_function,
        )
    raise RuntimeError(  # noqa: TRY003
        "Test context environment variables not set"  # noqa: EM101
    )


_codeflash_call_site: contextvars.ContextVar[str] = contextvars.ContextVar(
    "codeflash_call_site", default=""
)

CREATE_TABLE_SQL = (
    "CREATE TABLE IF NOT EXISTS codeflash_results ("
    "test_module_path TEXT NOT NULL, "
    "test_class_name TEXT, "
    "test_function_name TEXT NOT NULL, "
    "function_getting_tested TEXT NOT NULL, "
    "loop_index INTEGER NOT NULL, "
    "invocation_id TEXT NOT NULL, "
    "mode TEXT NOT NULL, "
    "wall_time_ns INTEGER NOT NULL, "
    "cpu_time_ns INTEGER, "
    "return_value BLOB, "
    "verification_type TEXT, "
    "sequential_time_ns INTEGER, "
    "concurrent_time_ns INTEGER, "
    "concurrency_factor INTEGER, "
    "stdout TEXT"
    ")"
)

connections: dict[str, sqlite3.Connection] = {}


def get_async_db(
    db_path: Path,
) -> tuple[sqlite3.Connection, sqlite3.Cursor]:
    """Get or create a cached SQLite connection."""
    key = str(db_path)
    if key not in connections:
        conn = sqlite3.connect(db_path)
        conn.execute(CREATE_TABLE_SQL)
        connections[key] = conn
    conn = connections[key]
    return conn, conn.cursor()


def close_all_connections() -> None:
    """Commit and close all cached connections."""
    for conn in connections.values():
        try:
            conn.commit()
            conn.close()
        except Exception:  # noqa: PERF203, S110
            pass
    connections.clear()


atexit.register(close_all_connections)


def detect_device_sync() -> tuple[
    Any | None,  # torch_cuda_sync
    Any | None,  # torch_mps_sync
    Any | None,  # tf_sync
    bool,  # jax_available
]:
    """Detect available GPU frameworks and return sync callables.

    Called once at first decorator invocation; results are cached.
    Uses ``find_spec`` to avoid importing heavy frameworks when absent.
    """
    from importlib.util import find_spec  # noqa: PLC0415

    torch_cuda_sync = None
    torch_mps_sync = None
    tf_sync = None
    jax_available = False

    if find_spec("torch") is not None:
        import torch  # type: ignore[import-not-found]  # noqa: PLC0415

        if torch.cuda.is_available() and torch.cuda.is_initialized():
            torch_cuda_sync = torch.cuda.synchronize
        elif (
            hasattr(torch, "backends")
            and hasattr(torch.backends, "mps")
            and torch.backends.mps.is_available()
            and hasattr(torch, "mps")
            and hasattr(torch.mps, "synchronize")
        ):
            torch_mps_sync = torch.mps.synchronize

    if find_spec("jax") is not None:
        import jax  # type: ignore[import-not-found]  # noqa: PLC0415

        jax_available = hasattr(jax, "block_until_ready")

    if find_spec("tensorflow") is not None:
        import tensorflow as tf  # type: ignore[import-untyped]  # noqa: PLC0415

        if hasattr(tf.test, "experimental") and hasattr(
            tf.test.experimental, "sync_devices"
        ):
            tf_sync = tf.test.experimental.sync_devices

    return torch_cuda_sync, torch_mps_sync, tf_sync, jax_available


device_sync_cache: tuple[Any | None, Any | None, Any | None, bool] | None = (
    None
)


def get_device_sync() -> tuple[Any | None, Any | None, Any | None, bool]:
    """Return cached device sync callables, detecting on first call."""
    global device_sync_cache  # noqa: PLW0603
    if device_sync_cache is None:
        device_sync_cache = detect_device_sync()
    return device_sync_cache


def sync_devices_before() -> None:
    """Synchronize GPU devices before timing."""
    cuda_sync, mps_sync, tf_sync, _ = get_device_sync()
    if cuda_sync is not None:
        cuda_sync()
    elif mps_sync is not None:
        mps_sync()
    if tf_sync is not None:
        tf_sync()


def sync_devices_after(return_value: Any) -> None:
    """Synchronize GPU devices after function call."""
    cuda_sync, mps_sync, tf_sync, jax_available = get_device_sync()
    if cuda_sync is not None:
        cuda_sync()
    elif mps_sync is not None:
        mps_sync()
    if jax_available and hasattr(return_value, "block_until_ready"):
        return_value.block_until_ready()
    if tf_sync is not None:
        tf_sync()


def codeflash_behavior_sync(func: F) -> F:
    """
    Capture sync return values, timing, and stdout for behavioral tests.

    Results are written to the codeflash_results SQLite table.
    """

    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        function_name = func.__name__
        call_site = _codeflash_call_site.get()
        loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
        (
            test_module_name,
            test_class_name,
            test_name,
        ) = extract_test_context_from_env()

        test_id = (
            f"{test_module_name}:{test_class_name}"
            f":{test_name}:{call_site}:{loop_index}"
        )

        if not hasattr(wrapper, "index"):
            wrapper.index = {}  # type: ignore[attr-defined]
        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]

        call_index = wrapper.index[test_id]  # type: ignore[attr-defined]
        invocation_id = f"{call_site}_{call_index}"

        iteration = os.environ.get("CODEFLASH_TEST_ITERATION", "0")
        db_path = get_run_tmp_file(
            Path(f"codeflash_results_{iteration}.sqlite")
        )
        conn, cur = get_async_db(db_path)

        exception = None
        captured_stdout = io.StringIO()
        old_stdout = sys.stdout
        sys.stdout = captured_stdout
        sync_devices_before()
        gc.disable()
        try:
            counter = time.perf_counter_ns()
            cpu_counter = time.thread_time_ns()
            return_value = func(*args, **kwargs)
            sync_devices_after(return_value)
            wall_time = time.perf_counter_ns() - counter
            cpu_time = time.thread_time_ns() - cpu_counter
        except Exception as e:
            wall_time = time.perf_counter_ns() - counter
            cpu_time = time.thread_time_ns() - cpu_counter
            exception = e
        finally:
            gc.enable()
            sys.stdout = old_stdout

        stdout_text = captured_stdout.getvalue()

        pickled = (
            pickle.dumps(exception)
            if exception
            else pickle.dumps((args, kwargs, return_value))
        )
        cur.execute(
            "INSERT INTO codeflash_results VALUES "
            "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                test_module_name,
                test_class_name,
                test_name,
                function_name,
                loop_index,
                invocation_id,
                "behavior",
                wall_time,
                cpu_time,
                pickled,
                VerificationType.FUNCTION_CALL.value,
                None,
                None,
                None,
                stdout_text,
            ),
        )
        conn.commit()

        if exception:
            raise exception
        return return_value

    return wrapper  # type: ignore[return-value]


def codeflash_performance_sync(func: F) -> F:
    """
    Measure sync execution time for performance tests.

    Results are written to the codeflash_results SQLite table.
    """

    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        function_name = func.__name__
        call_site = _codeflash_call_site.get()
        loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
        (
            test_module_name,
            test_class_name,
            test_name,
        ) = extract_test_context_from_env()

        test_id = (
            f"{test_module_name}:{test_class_name}"
            f":{test_name}:{call_site}:{loop_index}"
        )

        if not hasattr(wrapper, "index"):
            wrapper.index = {}  # type: ignore[attr-defined]
        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]

        call_index = wrapper.index[test_id]  # type: ignore[attr-defined]
        invocation_id = f"{call_site}_{call_index}"

        iteration = os.environ.get("CODEFLASH_TEST_ITERATION", "0")
        db_path = get_run_tmp_file(
            Path(f"codeflash_results_{iteration}.sqlite")
        )
        conn, cur = get_async_db(db_path)

        exception = None
        sync_devices_before()
        gc.disable()
        try:
            counter = time.perf_counter_ns()
            return_value = func(*args, **kwargs)
            sync_devices_after(return_value)
            wall_time = time.perf_counter_ns() - counter
        except Exception as e:
            wall_time = time.perf_counter_ns() - counter
            exception = e
        finally:
            gc.enable()

        cur.execute(
            "INSERT INTO codeflash_results VALUES "
            "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                test_module_name,
                test_class_name,
                test_name,
                function_name,
                loop_index,
                invocation_id,
                "performance",
                wall_time,
                None,
                None,
                None,
                None,
                None,
                None,
                None,
            ),
        )
        conn.commit()

        if exception:
            raise exception
        return return_value

    return wrapper  # type: ignore[return-value]


def codeflash_behavior_async(func: F) -> F:
    """
    Capture async return values and timing for behavioral tests.

    Results are written to the codeflash_results SQLite table.
    """

    @wraps(func)
    async def wrapper(*args: Any, **kwargs: Any) -> Any:
        loop = asyncio.get_running_loop()
        function_name = func.__name__
        call_site = _codeflash_call_site.get()
        loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
        (
            test_module_name,
            test_class_name,
            test_name,
        ) = extract_test_context_from_env()

        test_id = (
            f"{test_module_name}:{test_class_name}"
            f":{test_name}:{call_site}:{loop_index}"
        )

        if not hasattr(wrapper, "index"):
            wrapper.index = {}  # type: ignore[attr-defined]
        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]

        call_index = wrapper.index[test_id]  # type: ignore[attr-defined]
        invocation_id = f"{call_site}_{call_index}"

        iteration = os.environ.get("CODEFLASH_TEST_ITERATION", "0")
        db_path = get_run_tmp_file(
            Path(f"codeflash_results_{iteration}.sqlite")
        )
        conn, cur = get_async_db(db_path)

        exception = None
        captured_stdout = io.StringIO()
        old_stdout = sys.stdout
        sys.stdout = captured_stdout
        counter = loop.time()
        gc.disable()
        try:
            ret = func(*args, **kwargs)
            counter = loop.time()
            return_value = await ret
            wall_time = int((loop.time() - counter) * 1_000_000_000)
        except Exception as e:
            wall_time = int((loop.time() - counter) * 1_000_000_000)
            exception = e
        finally:
            gc.enable()
            sys.stdout = old_stdout

        stdout_text = captured_stdout.getvalue()

        pickled = (
            pickle.dumps(exception)
            if exception
            else pickle.dumps((args, kwargs, return_value))
        )
        cur.execute(
            "INSERT INTO codeflash_results VALUES "
            "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                test_module_name,
                test_class_name,
                test_name,
                function_name,
                loop_index,
                invocation_id,
                "behavior",
                wall_time,
                None,
                pickled,
                VerificationType.FUNCTION_CALL.value,
                None,
                None,
                None,
                stdout_text,
            ),
        )
        conn.commit()

        if exception:
            raise exception
        return return_value

    return wrapper  # type: ignore[return-value]


def codeflash_performance_async(func: F) -> F:
    """
    Measure async execution time for performance tests.

    Results are written to the codeflash_results SQLite table.
    """

    @wraps(func)
    async def wrapper(*args: Any, **kwargs: Any) -> Any:
        loop = asyncio.get_running_loop()
        call_site = _codeflash_call_site.get()
        loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
        (
            test_module_name,
            test_class_name,
            test_name,
        ) = extract_test_context_from_env()

        test_id = (
            f"{test_module_name}:{test_class_name}"
            f":{test_name}:{call_site}:{loop_index}"
        )

        if not hasattr(wrapper, "index"):
            wrapper.index = {}  # type: ignore[attr-defined]
        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]

        call_index = wrapper.index[test_id]  # type: ignore[attr-defined]
        invocation_id = f"{call_site}_{call_index}"

        iteration = os.environ.get("CODEFLASH_TEST_ITERATION", "0")
        db_path = get_run_tmp_file(
            Path(f"codeflash_results_{iteration}.sqlite")
        )
        conn, cur = get_async_db(db_path)

        exception = None
        counter = loop.time()
        gc.disable()
        try:
            ret = func(*args, **kwargs)
            counter = loop.time()
            return_value = await ret
            wall_time = int((loop.time() - counter) * 1_000_000_000)
        except Exception as e:
            wall_time = int((loop.time() - counter) * 1_000_000_000)
            exception = e
        finally:
            gc.enable()

        cur.execute(
            "INSERT INTO codeflash_results VALUES "
            "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                test_module_name,
                test_class_name,
                test_name,
                func.__name__,
                loop_index,
                invocation_id,
                "performance",
                wall_time,
                None,
                None,
                None,
                None,
                None,
                None,
                None,
            ),
        )
        conn.commit()

        if exception:
            raise exception
        return return_value

    return wrapper  # type: ignore[return-value]


def codeflash_concurrency_async(func: F) -> F:
    """
    Measure concurrent vs sequential execution for async functions.

    Results are written to the codeflash_results SQLite table.
    """

    @wraps(func)
    async def wrapper(*args: Any, **kwargs: Any) -> Any:
        function_name = func.__name__
        concurrency_factor = int(
            os.environ.get("CODEFLASH_CONCURRENCY_FACTOR", "10")
        )
        (
            test_module_name,
            test_class_name,
            test_name,
        ) = extract_test_context_from_env()
        loop_index = int(os.environ.get("CODEFLASH_LOOP_INDEX", "0"))

        iteration = os.environ.get("CODEFLASH_TEST_ITERATION", "0")
        db_path = get_run_tmp_file(
            Path(f"codeflash_results_{iteration}.sqlite")
        )
        conn, cur = get_async_db(db_path)

        gc.disable()
        try:
            seq_start = time.perf_counter_ns()
            for _ in range(concurrency_factor):
                result = await func(*args, **kwargs)
            sequential_time = time.perf_counter_ns() - seq_start
        finally:
            gc.enable()

        gc.disable()
        try:
            conc_start = time.perf_counter_ns()
            tasks = [func(*args, **kwargs) for _ in range(concurrency_factor)]
            await asyncio.gather(*tasks)
            concurrent_time = time.perf_counter_ns() - conc_start
        finally:
            gc.enable()

        cur.execute(
            "INSERT INTO codeflash_results VALUES "
            "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                test_module_name,
                test_class_name,
                test_name,
                function_name,
                loop_index,
                "",
                "concurrency",
                0,
                None,
                None,
                None,
                sequential_time,
                concurrent_time,
                concurrency_factor,
                None,
            ),
        )
        conn.commit()

        return result

    return wrapper  # type: ignore[return-value]


__all__ = [
    "VerificationType",
    "_codeflash_call_site",
    "codeflash_behavior_async",
    "codeflash_behavior_sync",
    "codeflash_concurrency_async",
    "codeflash_performance_async",
    "codeflash_performance_sync",
    "extract_test_context_from_env",
    "get_run_tmp_file",
]
