"""
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 os
import sqlite3
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."""
    if not hasattr(get_run_tmp_file, "tmpdir"):
        get_run_tmp_file.tmpdir = TemporaryDirectory(  # type: ignore[attr-defined]
            prefix="codeflash_"
        )
    return Path(get_run_tmp_file.tmpdir.name) / 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 async_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, "
    "return_value BLOB, "
    "verification_type TEXT, "
    "sequential_time_ns INTEGER, "
    "concurrent_time_ns INTEGER, "
    "concurrency_factor INTEGER"
    ")"
)

_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 codeflash_behavior_async(func: F) -> F:
    """
    Capture async return values and timing for behavioral tests.

    Results are written to the async_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"async_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()

        pickled = (
            pickle.dumps(exception)
            if exception
            else pickle.dumps((args, kwargs, return_value))
        )
        cur.execute(
            "INSERT INTO async_results VALUES "
            "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                test_module_name,
                test_class_name,
                test_name,
                function_name,
                loop_index,
                invocation_id,
                "behavior",
                wall_time,
                pickled,
                VerificationType.FUNCTION_CALL.value,
                None,
                None,
                None,
            ),
        )
        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 async_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"async_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 async_results VALUES "
            "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                test_module_name,
                test_class_name,
                test_name,
                func.__name__,
                loop_index,
                invocation_id,
                "performance",
                wall_time,
                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 async_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"async_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 async_results VALUES "
            "(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                test_module_name,
                test_class_name,
                test_name,
                function_name,
                loop_index,
                "",
                "concurrency",
                0,
                None,
                None,
                sequential_time,
                concurrent_time,
                concurrency_factor,
            ),
        )
        conn.commit()

        return result

    return wrapper  # type: ignore[return-value]


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