"""Unit tests for inject_profiling_into_existing_test with different used_frameworks values.

Tests verify that:
- ``detect_frameworks_from_code`` correctly identifies GPU framework imports
- The sync instrumentation path produces framework-agnostic output with
  ``_codeflash_call_site`` tracking instead of ``codeflash_wrap``
"""

from __future__ import annotations

from pathlib import Path

from codeflash_python._model import FunctionToOptimize, TestingMode
from codeflash_python.test_discovery.models import CodePosition
from codeflash_python.testing._instrument_core import (
    detect_frameworks_from_code,
)
from codeflash_python.testing._instrumentation import (
    inject_profiling_into_existing_test,
)


class TestDetectFrameworksFromCode:
    """Tests for the detect_frameworks_from_code helper function."""

    def test_no_frameworks(self) -> None:
        """Test detection with no GPU framework imports."""
        code = """import os
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {}
        assert result == expected

    def test_torch_import(self) -> None:
        """Test detection with torch import."""
        code = """import torch
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"torch": "torch"}
        assert result == expected

    def test_torch_aliased_import(self) -> None:
        """Test detection with torch imported as alias."""
        code = """import torch as th
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"torch": "th"}
        assert result == expected

    def test_torch_submodule_import(self) -> None:
        """Test detection with torch submodule import (from torch import nn)."""
        code = """from torch import nn
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"torch": "torch"}
        assert result == expected

    def test_torch_dotted_import(self) -> None:
        """Test detection with torch.cuda or torch.nn import."""
        code = """import torch.cuda
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"torch": "torch"}
        assert result == expected

    def test_tensorflow_import(self) -> None:
        """Test detection with tensorflow import."""
        code = """import tensorflow
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"tensorflow": "tensorflow"}
        assert result == expected

    def test_tensorflow_aliased_import(self) -> None:
        """Test detection with tensorflow imported as alias."""
        code = """import tensorflow as tf
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"tensorflow": "tf"}
        assert result == expected

    def test_tensorflow_submodule_import(self) -> None:
        """Test detection with tensorflow submodule import."""
        code = """from tensorflow import keras
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"tensorflow": "tensorflow"}
        assert result == expected

    def test_jax_import(self) -> None:
        """Test detection with jax import."""
        code = """import jax
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"jax": "jax"}
        assert result == expected

    def test_jax_aliased_import(self) -> None:
        """Test detection with jax imported as alias."""
        code = """import jax as jnp
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"jax": "jnp"}
        assert result == expected

    def test_jax_submodule_import(self) -> None:
        """Test detection with jax submodule import."""
        code = """from jax import numpy as jnp
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"jax": "jax"}
        assert result == expected

    def test_multiple_frameworks(self) -> None:
        """Test detection with multiple framework imports."""
        code = """import torch
import tensorflow
import jax
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"torch": "torch", "tensorflow": "tensorflow", "jax": "jax"}
        assert result == expected

    def test_multiple_frameworks_aliased(self) -> None:
        """Test detection with multiple aliased framework imports."""
        code = """import torch as th
import tensorflow as tf
import jax as jnp
from mymodule import my_function

def test_something():
    pass
"""
        result = detect_frameworks_from_code(code)
        expected = {"torch": "th", "tensorflow": "tf", "jax": "jnp"}
        assert result == expected

    def test_syntax_error_returns_empty(self) -> None:
        """Test that syntax errors return empty dict."""
        code = """this is not valid python code !!!"""
        result = detect_frameworks_from_code(code)
        expected = {}
        assert result == expected


def _make_func() -> FunctionToOptimize:
    """Create a FunctionToOptimize for test helpers."""
    return FunctionToOptimize(
        function_name="my_function",
        parents=[],
        file_path=Path("mymodule.py"),
    )


def _assert_sync_call_site_output(instrumented_code: str) -> None:
    """Assert the sync path output has call-site tracking, not codeflash_wrap."""
    assert (
        "from codeflash_async_wrapper import _codeflash_call_site"
        in instrumented_code
    )
    assert "_codeflash_call_site.set(" in instrumented_code
    assert "codeflash_wrap" not in instrumented_code
    assert "torch.cuda.synchronize" not in instrumented_code
    assert "torch.mps.synchronize" not in instrumented_code
    assert "tensorflow.test.experimental.sync_devices" not in instrumented_code
    assert "jax.block_until_ready" not in instrumented_code
    assert "import gc" not in instrumented_code
    assert "import sqlite3" not in instrumented_code
    assert "import dill" not in instrumented_code


class TestInjectProfilingBehaviorMode:
    """Tests for inject_profiling_into_existing_test in BEHAVIOR mode.

    The sync path produces framework-agnostic output with _codeflash_call_site
    tracking regardless of which GPU frameworks are imported.
    """

    def test_no_frameworks(self, tmp_path: Path) -> None:
        """Sync path injects call-site tracking with no frameworks present."""
        code = """from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        test_file = tmp_path / "test_example.py"
        test_file.write_text(code)

        success, instrumented_code = inject_profiling_into_existing_test(
            test_path=test_file,
            call_positions=[CodePosition(4, 13)],
            function_to_optimize=_make_func(),
            tests_project_root=tmp_path,
            mode=TestingMode.BEHAVIOR,
        )

        assert success is True
        _assert_sync_call_site_output(instrumented_code)

    def test_torch_import(self, tmp_path: Path) -> None:
        """Sync path is framework-agnostic even with torch import."""
        code = """import torch
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        test_file = tmp_path / "test_example.py"
        test_file.write_text(code)

        success, instrumented_code = inject_profiling_into_existing_test(
            test_path=test_file,
            call_positions=[CodePosition(5, 13)],
            function_to_optimize=_make_func(),
            tests_project_root=tmp_path,
            mode=TestingMode.BEHAVIOR,
        )

        assert success is True
        _assert_sync_call_site_output(instrumented_code)

    def test_tensorflow_import(self, tmp_path: Path) -> None:
        """Sync path is framework-agnostic even with tensorflow import."""
        code = """import tensorflow
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        test_file = tmp_path / "test_example.py"
        test_file.write_text(code)

        success, instrumented_code = inject_profiling_into_existing_test(
            test_path=test_file,
            call_positions=[CodePosition(5, 13)],
            function_to_optimize=_make_func(),
            tests_project_root=tmp_path,
            mode=TestingMode.BEHAVIOR,
        )

        assert success is True
        _assert_sync_call_site_output(instrumented_code)

    def test_all_frameworks(self, tmp_path: Path) -> None:
        """Sync path is framework-agnostic even with all GPU frameworks."""
        code = """import torch
import tensorflow
import jax
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        test_file = tmp_path / "test_example.py"
        test_file.write_text(code)

        success, instrumented_code = inject_profiling_into_existing_test(
            test_path=test_file,
            call_positions=[CodePosition(7, 13)],
            function_to_optimize=_make_func(),
            tests_project_root=tmp_path,
            mode=TestingMode.BEHAVIOR,
        )

        assert success is True
        _assert_sync_call_site_output(instrumented_code)


class TestInjectProfilingPerformanceMode:
    """Tests for inject_profiling_into_existing_test in PERFORMANCE mode.

    The sync path produces identical framework-agnostic output regardless of
    mode -- both BEHAVIOR and PERFORMANCE use _codeflash_call_site tracking.
    """

    def test_no_frameworks(self, tmp_path: Path) -> None:
        """Sync path injects call-site tracking with no frameworks present."""
        code = """from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        test_file = tmp_path / "test_example.py"
        test_file.write_text(code)

        success, instrumented_code = inject_profiling_into_existing_test(
            test_path=test_file,
            call_positions=[CodePosition(4, 13)],
            function_to_optimize=_make_func(),
            tests_project_root=tmp_path,
            mode=TestingMode.PERFORMANCE,
        )

        assert success is True
        _assert_sync_call_site_output(instrumented_code)

    def test_torch_import(self, tmp_path: Path) -> None:
        """Sync path is framework-agnostic even with torch import."""
        code = """import torch
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        test_file = tmp_path / "test_example.py"
        test_file.write_text(code)

        success, instrumented_code = inject_profiling_into_existing_test(
            test_path=test_file,
            call_positions=[CodePosition(5, 13)],
            function_to_optimize=_make_func(),
            tests_project_root=tmp_path,
            mode=TestingMode.PERFORMANCE,
        )

        assert success is True
        _assert_sync_call_site_output(instrumented_code)

    def test_tensorflow_import(self, tmp_path: Path) -> None:
        """Sync path is framework-agnostic even with tensorflow import."""
        code = """import tensorflow
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        test_file = tmp_path / "test_example.py"
        test_file.write_text(code)

        success, instrumented_code = inject_profiling_into_existing_test(
            test_path=test_file,
            call_positions=[CodePosition(5, 13)],
            function_to_optimize=_make_func(),
            tests_project_root=tmp_path,
            mode=TestingMode.PERFORMANCE,
        )

        assert success is True
        _assert_sync_call_site_output(instrumented_code)

    def test_all_frameworks(self, tmp_path: Path) -> None:
        """Sync path is framework-agnostic even with all GPU frameworks."""
        code = """import torch
import tensorflow
import jax
from mymodule import my_function

def test_my_function():
    result = my_function(1, 2)
    assert result == 3
"""
        test_file = tmp_path / "test_example.py"
        test_file.write_text(code)

        success, instrumented_code = inject_profiling_into_existing_test(
            test_path=test_file,
            call_positions=[CodePosition(7, 13)],
            function_to_optimize=_make_func(),
            tests_project_root=tmp_path,
            mode=TestingMode.PERFORMANCE,
        )

        assert success is True
        _assert_sync_call_site_output(instrumented_code)