"""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

import codeflash_python.runtime._codeflash_async_decorators as _deco_mod
from codeflash_python.runtime._codeflash_async_decorators import (
    VerificationType,
    _codeflash_call_site,
    close_all_connections,
    codeflash_behavior_async,
    codeflash_behavior_sync,
    codeflash_concurrency_async,
    codeflash_performance_async,
    codeflash_performance_sync,
    connections,
    detect_device_sync,
    extract_test_context_from_env,
    get_async_db,
    get_device_sync,
    get_run_tmp_file,
    sync_devices_after,
    sync_devices_before,
)
from conftest import has_module


@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="results_db_path")
def _results_db_path(env_setup):
    """Return the path where the results DB will be written."""
    iteration = env_setup["CODEFLASH_TEST_ITERATION"]
    db_path = get_run_tmp_file(Path(f"codeflash_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 codeflash_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='codeflash_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 testclose_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, results_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, results_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 results_db_path.exists()
        con = sqlite3.connect(results_db_path)
        cur = con.cursor()
        cur.execute("SELECT * FROM codeflash_results")
        rows = cur.fetchall()
        assert 1 == len(rows)
        row = rows[0]
        assert "behavior" == row[6]
        assert 0 < row[7]

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

    @pytest.mark.asyncio
    async def test_exception_handling(
        self, env_setup, results_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(results_db_path)
        cur = con.cursor()
        cur.execute("SELECT return_value FROM codeflash_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, results_db_path, capsys
    ) -> None:
        """Behavior decorator does not leak stdout to outer scope."""

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

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

    @pytest.mark.asyncio
    async def test_captures_stdout_in_sqlite(
        self, env_setup, results_db_path
    ) -> None:
        """Behavior decorator captures print output into the stdout column."""

        @codeflash_behavior_async
        async def greeter(name: str) -> str:
            print(f"hello {name}")
            return f"hi {name}"

        _codeflash_call_site.set("0")
        await greeter("world")
        close_all_connections()

        con = sqlite3.connect(results_db_path)
        cur = con.cursor()
        cur.execute("SELECT stdout FROM codeflash_results")
        row = cur.fetchone()
        assert "hello world\n" == row[0]
        con.close()


class TestBehaviorSync:
    """codeflash_behavior_sync decorator."""

    def test_returns_correct_value(self, env_setup, results_db_path) -> None:
        """Decorated function returns the original return value."""

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

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

    def test_writes_to_sqlite(self, env_setup, results_db_path) -> None:
        """Writes behavior result to async_results table."""

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

        _codeflash_call_site.set("0")
        multiply(5, 6)
        close_all_connections()

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

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

    def test_exception_handling(self, env_setup, results_db_path) -> None:
        """Re-raises exceptions and stores them pickled."""

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

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

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

    def test_captures_stdout_in_sqlite(
        self, env_setup, results_db_path
    ) -> None:
        """Captures print output into the stdout column."""

        @codeflash_behavior_sync
        def greeter(name: str) -> str:
            print(f"hello {name}")
            return f"hi {name}"

        _codeflash_call_site.set("0")
        greeter("world")
        close_all_connections()

        con = sqlite3.connect(results_db_path)
        cur = con.cursor()
        cur.execute("SELECT stdout FROM codeflash_results")
        row = cur.fetchone()
        assert "hello world\n" == row[0]
        con.close()

    def test_no_stdout_leak(self, env_setup, results_db_path, capsys) -> None:
        """Sync decorator does not leak stdout to outer scope."""

        @codeflash_behavior_sync
        def quiet() -> int:
            print("captured")
            return 1

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

    def test_records_cpu_time(self, env_setup, results_db_path) -> None:
        """Records cpu_time_ns in its own column."""

        @codeflash_behavior_sync
        def work() -> int:
            return sum(range(1000))

        _codeflash_call_site.set("0")
        work()
        close_all_connections()

        con = sqlite3.connect(results_db_path)
        cur = con.cursor()
        cur.execute("SELECT cpu_time_ns FROM codeflash_results")
        row = cur.fetchone()
        assert row[0] is not None
        assert 0 <= row[0]
        con.close()


class TestPerformanceSync:
    """codeflash_performance_sync decorator."""

    def test_returns_correct_value(self, env_setup, results_db_path) -> None:
        """Decorated function returns the original return value."""

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

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

    def test_writes_to_sqlite(self, env_setup, results_db_path) -> None:
        """Writes performance result with mode='performance' and null return_value."""

        @codeflash_performance_sync
        def work() -> int:
            return 42

        _codeflash_call_site.set("0")
        work()
        close_all_connections()

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

    def test_exception_handling(self, env_setup, results_db_path) -> None:
        """Re-raises exceptions from the wrapped function."""

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

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

    def test_no_stdout_capture(
        self, env_setup, results_db_path, capsys
    ) -> None:
        """Performance decorator does not redirect stdout."""

        @codeflash_performance_sync
        def talker() -> int:
            print("visible")
            return 1

        _codeflash_call_site.set("0")
        talker()
        captured = capsys.readouterr()
        assert "visible" in captured.out

    def test_records_wall_time(self, env_setup, results_db_path) -> None:
        """Records a positive wall_time_ns value."""

        @codeflash_performance_sync
        def work() -> int:
            return sum(range(1000))

        _codeflash_call_site.set("0")
        work()
        close_all_connections()

        con = sqlite3.connect(results_db_path)
        cur = con.cursor()
        cur.execute("SELECT wall_time_ns FROM codeflash_results")
        row = cur.fetchone()
        assert row[0] is not None
        assert 0 < row[0]
        con.close()


@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, results_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, results_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(results_db_path)
        cur = con.cursor()
        cur.execute("SELECT * FROM codeflash_results")
        rows = cur.fetchall()
        assert 1 == len(rows)
        row = rows[0]
        assert "performance" == row[6]
        assert 0 < row[7]
        assert row[9] is None
        assert row[10] is None
        con.close()

    @pytest.mark.asyncio
    async def test_no_stdout_output(
        self, env_setup, results_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, results_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, results_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(results_db_path)
        cur = con.cursor()
        cur.execute("SELECT * FROM codeflash_results")
        rows = cur.fetchall()
        assert 1 == len(rows)
        row = rows[0]
        assert "concurrency" == row[6]
        assert 0 == row[7]
        assert 0 < row[11]
        assert 0 < row[12]
        assert 3 == row[13]
        con.close()

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

    @pytest.mark.asyncio
    async def test_no_stdout_output(
        self, env_setup, results_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, results_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(results_db_path)
        cur = con.cursor()
        cur.execute(
            "SELECT invocation_id FROM codeflash_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, results_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, results_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(results_db_path)
        cur = con.cursor()
        cur.execute(
            "SELECT invocation_id FROM codeflash_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 codeflash_results SQLite schema."""

    def test_table_columns(self, tmp_path) -> None:
        """codeflash_results table has exactly 15 columns."""
        db_path = tmp_path / "schema_test.sqlite"
        conn, cur = get_async_db(db_path)
        cur.execute("PRAGMA table_info(codeflash_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",
            "cpu_time_ns",
            "return_value",
            "verification_type",
            "sequential_time_ns",
            "concurrent_time_ns",
            "concurrency_factor",
            "stdout",
        ]
        actual_names = [col[1] for col in columns]
        assert expected_names == actual_names
        conn.close()
        connections.pop(str(db_path), None)


@pytest.fixture(name="reset_device_cache")
def _reset_device_cache():
    """Reset the module-level device sync cache before and after each test."""
    _deco_mod.device_sync_cache = None
    yield
    _deco_mod.device_sync_cache = None


class TestDetectDeviceSync:
    """detect_device_sync probes real GPU frameworks via find_spec."""

    def test_returns_four_element_tuple(self, reset_device_cache) -> None:
        """Returns (cuda_sync, mps_sync, tf_sync, jax_available)."""
        result = detect_device_sync()
        assert 4 == len(result)

    @pytest.mark.skipif(not has_module("torch"), reason="torch not installed")
    def test_detects_real_torch(self, reset_device_cache) -> None:
        """Detects torch and returns a sync callable for the active device."""
        import torch

        cuda_sync, mps_sync, _, _ = detect_device_sync()
        if torch.cuda.is_available() and torch.cuda.is_initialized():
            assert cuda_sync is torch.cuda.synchronize
            assert mps_sync is None
        elif (
            hasattr(torch.backends, "mps")
            and torch.backends.mps.is_available()
        ):
            assert cuda_sync is None
            assert mps_sync is torch.mps.synchronize
        else:
            assert cuda_sync is None
            assert mps_sync is None

    @pytest.mark.skipif(not has_module("jax"), reason="jax not installed")
    def test_detects_real_jax(self, reset_device_cache) -> None:
        """Detects JAX and sets jax_available based on block_until_ready."""
        import jax

        _, _, _, jax_avail = detect_device_sync()
        assert jax_avail is hasattr(jax, "block_until_ready")

    @pytest.mark.skipif(
        not has_module("tensorflow"), reason="tensorflow not installed"
    )
    def test_detects_real_tensorflow(self, reset_device_cache) -> None:
        """Detects TensorFlow sync_devices when available."""
        import tensorflow as tf

        _, _, tf_sync, _ = detect_device_sync()
        if hasattr(tf.test, "experimental") and hasattr(
            tf.test.experimental, "sync_devices"
        ):
            assert tf_sync is tf.test.experimental.sync_devices
        else:
            assert tf_sync is None


class TestGetDeviceSync:
    """get_device_sync caches detect_device_sync results."""

    def test_caches_result(self, reset_device_cache) -> None:
        """Returns the same tuple on repeated calls without re-detecting."""
        first = get_device_sync()
        second = get_device_sync()
        assert first is second

    def test_redetects_after_cache_clear(self, reset_device_cache) -> None:
        """Re-runs detection after the cache is cleared."""
        first = get_device_sync()
        _deco_mod.device_sync_cache = None
        second = get_device_sync()
        assert first == second
        assert first is not second


class TestSyncDevicesBefore:
    """sync_devices_before exercises real framework sync paths."""

    def test_runs_without_error(self, reset_device_cache) -> None:
        """Calling with real frameworks installed does not raise."""
        sync_devices_before()

    def test_cuda_takes_priority_over_mps(self, reset_device_cache) -> None:
        """CUDA sync is called instead of MPS when both are in the cache."""
        calls = []
        _deco_mod.device_sync_cache = (
            lambda: calls.append("cuda"),
            lambda: calls.append("mps"),
            None,
            False,
        )
        sync_devices_before()
        assert ["cuda"] == calls

    def test_mps_called_when_no_cuda(self, reset_device_cache) -> None:
        """MPS sync fires when CUDA is absent in the cache."""
        calls = []
        _deco_mod.device_sync_cache = (
            None,
            lambda: calls.append("mps"),
            None,
            False,
        )
        sync_devices_before()
        assert ["mps"] == calls

    def test_tf_called_independently(self, reset_device_cache) -> None:
        """TF sync fires independently of torch sync."""
        calls = []
        _deco_mod.device_sync_cache = (
            None,
            None,
            lambda: calls.append("tf"),
            False,
        )
        sync_devices_before()
        assert ["tf"] == calls


class TestSyncDevicesAfter:
    """sync_devices_after exercises real framework sync paths."""

    def test_runs_without_error(self, reset_device_cache) -> None:
        """Calling with real frameworks and a plain return value does not raise."""
        sync_devices_after(42)

    @pytest.mark.skipif(not has_module("jax"), reason="jax not installed")
    def test_jax_block_until_ready_on_real_array(
        self, reset_device_cache
    ) -> None:
        """Calls block_until_ready on a real JAX array."""
        import jax.numpy as jnp

        arr = jnp.array([1, 2, 3])
        sync_devices_after(arr)

    def test_skips_jax_on_plain_value(self, reset_device_cache) -> None:
        """Does not fail when jax_available=True but return value is plain."""
        _deco_mod.device_sync_cache = (None, None, None, True)
        sync_devices_after(42)

    def test_cuda_priority_in_after(self, reset_device_cache) -> None:
        """CUDA sync fires instead of MPS in the after path too."""
        calls = []
        _deco_mod.device_sync_cache = (
            lambda: calls.append("cuda"),
            lambda: calls.append("mps"),
            None,
            False,
        )
        sync_devices_after(42)
        assert ["cuda"] == calls

    @pytest.mark.skipif(not has_module("jax"), reason="jax not installed")
    def test_all_syncs_fire_together(self, reset_device_cache) -> None:
        """All applicable syncs fire: torch + JAX block_until_ready + TF."""
        import jax.numpy as jnp

        calls = []
        _deco_mod.device_sync_cache = (
            lambda: calls.append("cuda"),
            None,
            lambda: calls.append("tf"),
            True,
        )
        arr = jnp.array([1.0])
        sync_devices_after(arr)
        assert ["cuda", "tf"] == calls
        assert hasattr(arr, "block_until_ready")
