"""Tests for the new self-contained async instrumentation decorators."""

from __future__ import annotations

import asyncio
import contextvars
import os
import sqlite3
import sys
from pathlib import Path
from unittest.mock import patch

import dill as pickle
import pytest

from codeflash_python.runtime._codeflash_async_decorators import (
    VerificationType,
    _close_all_connections,
    _codeflash_call_site,
    _connections,
    _get_async_db,
    codeflash_behavior_async,
    codeflash_concurrency_async,
    codeflash_performance_async,
    extract_test_context_from_env,
    get_run_tmp_file,
)


@pytest.fixture(name="env_setup")
def _env_setup(request, tmp_path):
    """Set up env vars and clean up after each test."""
    original_env = {}
    test_env = {
        "CODEFLASH_LOOP_INDEX": "1",
        "CODEFLASH_TEST_ITERATION": "0",
        "CODEFLASH_TEST_MODULE": __name__,
        "CODEFLASH_TEST_CLASS": "",
        "CODEFLASH_TEST_FUNCTION": request.node.name,
    }
    for key, value in test_env.items():
        original_env[key] = os.environ.get(key)
        os.environ[key] = value

    yield test_env

    for key, original_value in original_env.items():
        if original_value is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = original_value

    _close_all_connections()


@pytest.fixture(name="async_db_path")
def _async_db_path(env_setup):
    """Return the path where the async results DB will be written."""
    iteration = env_setup["CODEFLASH_TEST_ITERATION"]
    db_path = get_run_tmp_file(Path(f"async_results_{iteration}.sqlite"))
    yield db_path
    if db_path.exists():
        db_path.unlink()


class TestExtractTestContextFromEnv:
    """extract_test_context_from_env reads CODEFLASH_TEST_* env vars."""

    def test_returns_tuple(self) -> None:
        """Returns (module, class_or_none, function) from env."""
        with patch.dict(
            os.environ,
            {
                "CODEFLASH_TEST_MODULE": "mod",
                "CODEFLASH_TEST_CLASS": "Cls",
                "CODEFLASH_TEST_FUNCTION": "test_fn",
            },
        ):
            module, cls, fn = extract_test_context_from_env()
        assert "mod" == module
        assert "Cls" == cls
        assert "test_fn" == fn

    def test_class_none_when_empty(self) -> None:
        """Returns None for test class when env var is empty."""
        with patch.dict(
            os.environ,
            {
                "CODEFLASH_TEST_MODULE": "mod",
                "CODEFLASH_TEST_CLASS": "",
                "CODEFLASH_TEST_FUNCTION": "test_fn",
            },
        ):
            _, cls, _ = extract_test_context_from_env()
        assert cls is None

    def test_raises_when_module_missing(self) -> None:
        """Raises KeyError when CODEFLASH_TEST_MODULE is unset."""
        env = os.environ.copy()
        env.pop("CODEFLASH_TEST_MODULE", None)
        with patch.dict(os.environ, env, clear=True), pytest.raises(KeyError):
            extract_test_context_from_env()


class TestCodeflashCallSite:
    """_codeflash_call_site contextvar behavior."""

    def test_default_value(self) -> None:
        """Default value is empty string."""
        ctx = contextvars.copy_context()
        assert "" == ctx.run(_codeflash_call_site.get)

    def test_set_and_get(self) -> None:
        """Can set and retrieve a value."""
        token = _codeflash_call_site.set("42")
        assert "42" == _codeflash_call_site.get()
        _codeflash_call_site.reset(token)


class TestGetAsyncDb:
    """_get_async_db connection caching."""

    def test_creates_table(self, tmp_path) -> None:
        """Creates the async_results table on first connect."""
        db_path = tmp_path / "test.sqlite"
        conn, cur = _get_async_db(db_path)
        cur.execute(
            "SELECT name FROM sqlite_master "
            "WHERE type='table' AND name='async_results'"
        )
        assert cur.fetchone() is not None
        conn.close()
        _connections.pop(str(db_path), None)

    def test_caches_connection(self, tmp_path) -> None:
        """Returns the same connection on repeated calls."""
        db_path = tmp_path / "test.sqlite"
        conn1, _ = _get_async_db(db_path)
        conn2, _ = _get_async_db(db_path)
        assert conn1 is conn2
        conn1.close()
        _connections.pop(str(db_path), None)

    def test_close_all_connections(self, tmp_path) -> None:
        """_close_all_connections empties the cache."""
        db_path = tmp_path / "test.sqlite"
        _get_async_db(db_path)
        assert 0 < len(_connections)
        _close_all_connections()
        assert 0 == len(_connections)


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="pending support for asyncio on windows",
)
class TestBehaviorAsync:
    """codeflash_behavior_async decorator."""

    @pytest.mark.asyncio
    async def test_returns_correct_value(
        self, env_setup, async_db_path
    ) -> None:
        """Decorated function returns the original return value."""

        @codeflash_behavior_async
        async def add(a: int, b: int) -> int:
            return a + b

        _codeflash_call_site.set("0")
        result = await add(3, 4)
        assert 7 == result

    @pytest.mark.asyncio
    async def test_writes_to_sqlite(self, env_setup, async_db_path) -> None:
        """Writes behavior result to async_results table."""

        @codeflash_behavior_async
        async def multiply(a: int, b: int) -> int:
            return a * b

        _codeflash_call_site.set("0")
        await multiply(5, 6)
        _close_all_connections()

        assert async_db_path.exists()
        con = sqlite3.connect(async_db_path)
        cur = con.cursor()
        cur.execute("SELECT * FROM async_results")
        rows = cur.fetchall()
        assert 1 == len(rows)
        row = rows[0]
        assert "behavior" == row[6]
        assert 0 < row[7]

        data = pickle.loads(row[8])
        args, kwargs, ret = data
        assert (5, 6) == args
        assert {} == kwargs
        assert 30 == ret
        assert VerificationType.FUNCTION_CALL.value == row[9]
        con.close()

    @pytest.mark.asyncio
    async def test_exception_handling(self, env_setup, async_db_path) -> None:
        """Re-raises exceptions and stores them pickled."""

        @codeflash_behavior_async
        async def fail() -> None:
            raise ValueError("boom")

        _codeflash_call_site.set("0")
        with pytest.raises(ValueError, match="boom"):
            await fail()

        _close_all_connections()
        con = sqlite3.connect(async_db_path)
        cur = con.cursor()
        cur.execute("SELECT return_value FROM async_results")
        row = cur.fetchone()
        exc = pickle.loads(row[0])
        assert isinstance(exc, ValueError)
        assert "boom" == str(exc)
        con.close()

    @pytest.mark.asyncio
    async def test_no_stdout_output(
        self, env_setup, async_db_path, capsys
    ) -> None:
        """Behavior decorator emits no stdout."""

        @codeflash_behavior_async
        async def noop() -> int:
            return 1

        _codeflash_call_site.set("0")
        await noop()
        captured = capsys.readouterr()
        assert "" == captured.out


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="pending support for asyncio on windows",
)
class TestPerformanceAsync:
    """codeflash_performance_async decorator."""

    @pytest.mark.asyncio
    async def test_returns_correct_value(
        self, env_setup, async_db_path
    ) -> None:
        """Returns the original return value."""

        @codeflash_performance_async
        async def add(a: int, b: int) -> int:
            return a + b

        _codeflash_call_site.set("0")
        result = await add(3, 4)
        assert 7 == result

    @pytest.mark.asyncio
    async def test_writes_to_sqlite(self, env_setup, async_db_path) -> None:
        """Writes performance result with null return_value."""

        @codeflash_performance_async
        async def work() -> int:
            return 42

        _codeflash_call_site.set("0")
        await work()
        _close_all_connections()

        con = sqlite3.connect(async_db_path)
        cur = con.cursor()
        cur.execute("SELECT * FROM async_results")
        rows = cur.fetchall()
        assert 1 == len(rows)
        row = rows[0]
        assert "performance" == row[6]
        assert 0 < row[7]
        assert row[8] is None
        assert row[9] is None
        con.close()

    @pytest.mark.asyncio
    async def test_no_stdout_output(
        self, env_setup, async_db_path, capsys
    ) -> None:
        """Performance decorator emits no stdout."""

        @codeflash_performance_async
        async def noop() -> int:
            return 1

        _codeflash_call_site.set("0")
        await noop()
        captured = capsys.readouterr()
        assert "" == captured.out


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="pending support for asyncio on windows",
)
class TestConcurrencyAsync:
    """codeflash_concurrency_async decorator."""

    @pytest.mark.asyncio
    async def test_returns_correct_value(
        self, env_setup, async_db_path
    ) -> None:
        """Returns the result from sequential execution."""

        @codeflash_concurrency_async
        async def add(a: int, b: int) -> int:
            return a + b

        result = await add(3, 4)
        assert 7 == result

    @pytest.mark.asyncio
    async def test_writes_concurrency_metrics(
        self, env_setup, async_db_path
    ) -> None:
        """Writes sequential/concurrent timing to async_results."""
        os.environ["CODEFLASH_CONCURRENCY_FACTOR"] = "3"

        @codeflash_concurrency_async
        async def work() -> int:
            await asyncio.sleep(0.001)
            return 42

        await work()
        _close_all_connections()

        con = sqlite3.connect(async_db_path)
        cur = con.cursor()
        cur.execute("SELECT * FROM async_results")
        rows = cur.fetchall()
        assert 1 == len(rows)
        row = rows[0]
        assert "concurrency" == row[6]
        assert 0 == row[7]
        assert 0 < row[10]
        assert 0 < row[11]
        assert 3 == row[12]
        con.close()

        os.environ.pop("CODEFLASH_CONCURRENCY_FACTOR", None)

    @pytest.mark.asyncio
    async def test_no_stdout_output(
        self, env_setup, async_db_path, capsys
    ) -> None:
        """Concurrency decorator emits no stdout."""
        os.environ["CODEFLASH_CONCURRENCY_FACTOR"] = "2"

        @codeflash_concurrency_async
        async def noop() -> int:
            return 1

        await noop()
        captured = capsys.readouterr()
        assert "" == captured.out

        os.environ.pop("CODEFLASH_CONCURRENCY_FACTOR", None)


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="pending support for asyncio on windows",
)
class TestBehaviorAsyncEdgeCases:
    """Edge cases for codeflash_behavior_async."""

    @pytest.mark.asyncio
    async def test_multiple_calls_increment_index(
        self, env_setup, async_db_path
    ) -> None:
        """Second call to same decorator increments the invocation counter."""

        @codeflash_behavior_async
        async def inc(x: int) -> int:
            return x + 1

        _codeflash_call_site.set("0")
        await inc(1)
        await inc(2)
        _close_all_connections()

        con = sqlite3.connect(async_db_path)
        cur = con.cursor()
        cur.execute(
            "SELECT invocation_id FROM async_results ORDER BY invocation_id"
        )
        ids = [row[0] for row in cur.fetchall()]
        assert 2 == len(ids)
        assert "0_0" == ids[0]
        assert "0_1" == ids[1]
        con.close()


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="pending support for asyncio on windows",
)
class TestPerformanceAsyncEdgeCases:
    """Edge cases for codeflash_performance_async."""

    @pytest.mark.asyncio
    async def test_exception_handling(self, env_setup, async_db_path) -> None:
        """Re-raises exceptions from the wrapped function."""

        @codeflash_performance_async
        async def fail() -> None:
            raise ValueError("perf boom")

        _codeflash_call_site.set("0")
        with pytest.raises(ValueError, match="perf boom"):
            await fail()

    @pytest.mark.asyncio
    async def test_multiple_calls_increment_index(
        self, env_setup, async_db_path
    ) -> None:
        """Second call increments the invocation counter."""

        @codeflash_performance_async
        async def work() -> int:
            return 1

        _codeflash_call_site.set("0")
        await work()
        await work()
        _close_all_connections()

        con = sqlite3.connect(async_db_path)
        cur = con.cursor()
        cur.execute(
            "SELECT invocation_id FROM async_results ORDER BY invocation_id"
        )
        ids = [row[0] for row in cur.fetchall()]
        assert "0_0" == ids[0]
        assert "0_1" == ids[1]
        con.close()


class TestCloseAllConnectionsErrorHandling:
    """_close_all_connections handles exceptions gracefully."""

    def test_handles_already_closed_connection(self, tmp_path) -> None:
        """Does not raise when a connection is already closed."""
        db_path = tmp_path / "test.sqlite"
        conn, _ = _get_async_db(db_path)
        conn.close()
        _close_all_connections()
        assert 0 == len(_connections)


class TestExtractTestContextEdgeCases:
    """Edge cases for extract_test_context_from_env."""

    def test_raises_when_module_empty(self) -> None:
        """Raises RuntimeError when module is empty string."""
        with (
            patch.dict(
                os.environ,
                {
                    "CODEFLASH_TEST_MODULE": "",
                    "CODEFLASH_TEST_CLASS": "",
                    "CODEFLASH_TEST_FUNCTION": "test_fn",
                },
            ),
            pytest.raises(RuntimeError),
        ):
            extract_test_context_from_env()


class TestSchemaValidation:
    """Validate the async_results SQLite schema."""

    def test_table_columns(self, tmp_path) -> None:
        """async_results table has exactly 13 columns."""
        db_path = tmp_path / "schema_test.sqlite"
        conn, cur = _get_async_db(db_path)
        cur.execute("PRAGMA table_info(async_results)")
        columns = cur.fetchall()
        expected_names = [
            "test_module_path",
            "test_class_name",
            "test_function_name",
            "function_getting_tested",
            "loop_index",
            "invocation_id",
            "mode",
            "wall_time_ns",
            "return_value",
            "verification_type",
            "sequential_time_ns",
            "concurrent_time_ns",
            "concurrency_factor",
        ]
        actual_names = [col[1] for col in columns]
        assert expected_names == actual_names
        conn.close()
        _connections.pop(str(db_path), None)
