"""Tests for line profiler instrumentation.

Ported from the original codeflash repo's
tests/test_instrument_line_profiler.py.
"""

from __future__ import annotations

import os
from pathlib import Path
from tempfile import TemporaryDirectory

from codeflash_python._model import FunctionToOptimize
from codeflash_python.benchmarking._line_profiling import add_decorator_imports
from codeflash_python.context.models import CodeOptimizationContext
from codeflash_python.context.pipeline import get_code_optimization_context
from codeflash_python.pipeline._function_optimizer import (
    write_code_and_helpers,
)
from codeflash_python.verification._baseline import contains_jit_decorator


def test_add_decorator_imports_helper_in_class() -> None:
    """Add line_profiler decorators to function with class-based helper."""
    code_path = (
        Path(__file__).parent.resolve()
        / "code_to_optimize/bubble_sort_classmethod.py"
    ).resolve()
    project_root_path = (Path(__file__).parent / "..").resolve()
    run_cwd = Path(__file__).parent.parent.resolve()
    func = FunctionToOptimize(
        function_name="sort_classmethod",
        parents=(),
        file_path=code_path,
    )
    original_cwd = os.getcwd()
    os.chdir(run_cwd)
    try:
        code_context: CodeOptimizationContext = get_code_optimization_context(
            func, project_root_path
        )
        original_helper_code: dict[Path, str] = {}
        helper_function_paths = {
            hf.file_path for hf in code_context.helper_functions
        }
        for helper_function_path in helper_function_paths:
            original_helper_code[helper_function_path] = (
                helper_function_path.read_text(encoding="utf8")
            )
        original_source = code_path.read_text(encoding="utf8")
        line_profiler_output_file = add_decorator_imports(
            func, code_context.helper_functions
        )
        expected_code_main = (
            "from line_profiler import profile as codeflash_line_profile\n"
            f"codeflash_line_profile.enable(output_prefix='{line_profiler_output_file.as_posix()}')\n"
            "from code_to_optimize.bubble_sort_in_class import BubbleSortClass\n"
            "\n\n@codeflash_line_profile\n"
            "def sort_classmethod(x):\n"
            "    y = BubbleSortClass()\n"
            "    return y.sorter(x)\n"
        )
        assert code_path.read_text("utf-8") == expected_code_main
        expected_code_helper = (
            "from line_profiler import profile as codeflash_line_profile\n"
            "def hi():\n    pass\n"
            "\n\nclass BubbleSortClass:\n"
            "    @codeflash_line_profile\n"
            "    def __init__(self):\n        pass\n"
            "\n    @codeflash_line_profile\n"
            "    def sorter(self, arr):\n"
            "        n = len(arr)\n"
            "        for i in range(n):\n"
            "            for j in range(n - i - 1):\n"
            "                if arr[j] > arr[j + 1]:\n"
            "                    arr[j], arr[j + 1] = arr[j + 1], arr[j]\n"
            "        return arr\n"
            "\n    def helper(self, arr, j):\n"
            "        return arr[j] > arr[j + 1]\n"
        )
        assert (
            code_context.helper_functions[0].file_path.read_text("utf-8")
            == expected_code_helper
        )
    finally:
        write_code_and_helpers(
            original_source, original_helper_code, func.file_path
        )
        os.chdir(original_cwd)


def test_add_decorator_imports_helper_in_nested_class() -> None:
    """Add line_profiler decorators to function with nested class helper."""
    code_path = (
        Path(__file__).parent.resolve()
        / "code_to_optimize/bubble_sort_nested_classmethod.py"
    ).resolve()
    project_root_path = (Path(__file__).parent / "..").resolve()
    run_cwd = Path(__file__).parent.parent.resolve()
    func = FunctionToOptimize(
        function_name="sort_classmethod",
        parents=(),
        file_path=code_path,
    )
    original_cwd = os.getcwd()
    os.chdir(run_cwd)
    try:
        code_context: CodeOptimizationContext = get_code_optimization_context(
            func, project_root_path
        )
        original_helper_code: dict[Path, str] = {}
        helper_function_paths = {
            hf.file_path for hf in code_context.helper_functions
        }
        for helper_function_path in helper_function_paths:
            original_helper_code[helper_function_path] = (
                helper_function_path.read_text(encoding="utf8")
            )
        original_source = code_path.read_text(encoding="utf8")
        line_profiler_output_file = add_decorator_imports(
            func, code_context.helper_functions
        )
        expected_code_main = (
            "from line_profiler import profile as codeflash_line_profile\n"
            f"codeflash_line_profile.enable(output_prefix='{line_profiler_output_file.as_posix()}')\n"
            "from code_to_optimize.bubble_sort_in_nested_class import WrapperClass\n"
            "\n\n@codeflash_line_profile\n"
            "def sort_classmethod(x):\n"
            "    y = WrapperClass.BubbleSortClass()\n"
            "    return y.sorter(x)\n"
        )
        assert code_path.read_text("utf-8") == expected_code_main
    finally:
        write_code_and_helpers(
            original_source, original_helper_code, func.file_path
        )
        os.chdir(original_cwd)


def test_add_decorator_imports_nodeps() -> None:
    """Add line_profiler decorators to function with no dependencies."""
    code_path = (
        Path(__file__).parent.resolve() / "code_to_optimize/bubble_sort.py"
    ).resolve()
    project_root_path = (Path(__file__).parent / "..").resolve()
    run_cwd = Path(__file__).parent.parent.resolve()
    func = FunctionToOptimize(
        function_name="sorter",
        parents=(),
        file_path=code_path,
    )
    original_cwd = os.getcwd()
    os.chdir(run_cwd)
    try:
        code_context: CodeOptimizationContext = get_code_optimization_context(
            func, project_root_path
        )
        original_helper_code: dict[Path, str] = {}
        helper_function_paths = {
            hf.file_path for hf in code_context.helper_functions
        }
        for helper_function_path in helper_function_paths:
            original_helper_code[helper_function_path] = (
                helper_function_path.read_text(encoding="utf8")
            )
        original_source = code_path.read_text(encoding="utf8")
        line_profiler_output_file = add_decorator_imports(
            func, code_context.helper_functions
        )
        expected_code_main = (
            "from line_profiler import profile as codeflash_line_profile\n"
            f"codeflash_line_profile.enable(output_prefix='{line_profiler_output_file.as_posix()}')\n"
            "@codeflash_line_profile\n"
            "def sorter(arr):\n"
            '    print("codeflash stdout: Sorting list")\n'
            "    for i in range(len(arr)):\n"
            "        for j in range(len(arr) - 1):\n"
            "            if arr[j] > arr[j + 1]:\n"
            "                temp = arr[j]\n"
            "                arr[j] = arr[j + 1]\n"
            "                arr[j + 1] = temp\n"
            '    print(f"result: {arr}")\n'
            "    return arr\n"
        )
        assert code_path.read_text("utf-8") == expected_code_main
    finally:
        write_code_and_helpers(
            original_source, original_helper_code, func.file_path
        )
        os.chdir(original_cwd)


def test_add_decorator_imports_helper_outside() -> None:
    """Add line_profiler decorators to function with helpers in other files."""
    code_path = (
        Path(__file__).parent.resolve()
        / "code_to_optimize/bubble_sort_deps.py"
    ).resolve()
    project_root_path = (Path(__file__).parent / "..").resolve()
    run_cwd = Path(__file__).parent.parent.resolve()
    func = FunctionToOptimize(
        function_name="sorter_deps",
        parents=(),
        file_path=code_path,
    )
    original_cwd = os.getcwd()
    os.chdir(run_cwd)
    try:
        code_context: CodeOptimizationContext = get_code_optimization_context(
            func, project_root_path
        )
        original_helper_code: dict[Path, str] = {}
        helper_function_paths = {
            hf.file_path for hf in code_context.helper_functions
        }
        for helper_function_path in helper_function_paths:
            original_helper_code[helper_function_path] = (
                helper_function_path.read_text(encoding="utf8")
            )
        original_source = code_path.read_text(encoding="utf8")
        line_profiler_output_file = add_decorator_imports(
            func, code_context.helper_functions
        )
        expected_code_main = (
            "from line_profiler import profile as codeflash_line_profile\n"
            f"codeflash_line_profile.enable(output_prefix='{line_profiler_output_file.as_posix()}')\n"
            "from code_to_optimize.bubble_sort_dep1_helper import dep1_comparer\n"
            "from code_to_optimize.bubble_sort_dep2_swap import dep2_swap\n"
            "\n\n@codeflash_line_profile\n"
            "def sorter_deps(arr):\n"
            "    for i in range(len(arr)):\n"
            "        for j in range(len(arr) - 1):\n"
            "            if dep1_comparer(arr, j):\n"
            "                dep2_swap(arr, j)\n"
            "    return arr\n\n"
        )
        expected_code_helper1 = (
            "from line_profiler import profile as codeflash_line_profile\n"
            "@codeflash_line_profile\n"
            "def dep1_comparer(arr, j: int) -> bool:\n"
            "    return arr[j] > arr[j + 1]\n"
        )
        expected_code_helper2 = (
            "from line_profiler import profile as codeflash_line_profile\n"
            "@codeflash_line_profile\n"
            "def dep2_swap(arr, j):\n"
            "    temp = arr[j]\n"
            "    arr[j] = arr[j + 1]\n"
            "    arr[j + 1] = temp\n"
        )
        assert code_path.read_text("utf-8") == expected_code_main
        assert (
            code_context.helper_functions[0].file_path.read_text("utf-8")
            == expected_code_helper1
        )
        assert (
            code_context.helper_functions[1].file_path.read_text("utf-8")
            == expected_code_helper2
        )
    finally:
        write_code_and_helpers(
            original_source, original_helper_code, func.file_path
        )
        os.chdir(original_cwd)


def test_add_decorator_imports_helper_in_dunder_class() -> None:
    """Add line_profiler decorators when helper is a class with __init__."""
    code_str = (
        "def sorter(arr):\n"
        "    ans = helper(arr)\n"
        "    return ans\n"
        "class helper:\n"
        "    def __init__(self, arr):\n"
        "        return arr.sort()"
    )
    code_path_dir = TemporaryDirectory()
    code_write_path = Path(code_path_dir.name) / "dunder_class.py"
    code_write_path.write_text(code_str, "utf-8")
    project_root_path = Path(code_path_dir.name)
    run_cwd = Path(__file__).parent.parent.resolve()
    func = FunctionToOptimize(
        function_name="sorter",
        parents=(),
        file_path=code_write_path,
    )
    original_cwd = os.getcwd()
    os.chdir(run_cwd)
    try:
        code_context: CodeOptimizationContext = get_code_optimization_context(
            func, project_root_path
        )
        line_profiler_output_file = add_decorator_imports(
            func, code_context.helper_functions
        )
        expected_code_main = (
            "from line_profiler import profile as codeflash_line_profile\n"
            f"codeflash_line_profile.enable(output_prefix='{line_profiler_output_file.as_posix()}')\n"
            "@codeflash_line_profile\n"
            "def sorter(arr):\n"
            "    ans = helper(arr)\n"
            "    return ans\n"
            "class helper:\n"
            "    @codeflash_line_profile\n"
            "    def __init__(self, arr):\n"
            "        return arr.sort()"
        )
        assert code_write_path.read_text("utf-8") == expected_code_main
    finally:
        os.chdir(original_cwd)


class TestContainsJitDecoratorNumba:
    """Tests for numba JIT decorator detection."""

    def test_numba_jit_with_module_prefix(self) -> None:
        """Detects @numba.jit with module prefix."""
        code = """
import numba

@numba.jit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_jit_with_alias(self) -> None:
        """Detects @nb.jit when numba is aliased."""
        code = """
import numba as nb

@nb.jit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_jit_direct_import(self) -> None:
        """Detects @jit via from-import."""
        code = """
from numba import jit

@jit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_jit_direct_import_with_alias(self) -> None:
        """Detects @my_jit when jit is aliased."""
        code = """
from numba import jit as my_jit

@my_jit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_jit_with_arguments(self) -> None:
        """Detects @numba.jit(nopython=True) with arguments."""
        code = """
import numba

@numba.jit(nopython=True)
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_jit_direct_import_with_arguments(self) -> None:
        """Detects @jit(nopython=True, cache=True) with arguments."""
        code = """
from numba import jit

@jit(nopython=True, cache=True)
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_njit(self) -> None:
        """Detects @njit via from-import."""
        code = """
from numba import njit

@njit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_njit_with_module_prefix(self) -> None:
        """Detects @numba.njit with module prefix."""
        code = """
import numba

@numba.njit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_vectorize(self) -> None:
        """Detects @vectorize via from-import."""
        code = """
from numba import vectorize

@vectorize
def my_func(x):
    return x * 2
"""
        assert contains_jit_decorator(code)

    def test_numba_guvectorize(self) -> None:
        """Detects @numba.guvectorize with arguments."""
        code = """
import numba

@numba.guvectorize(['void(float64[:], float64[:])'], '(n)->(n)')
def my_func(x, res):
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_stencil(self) -> None:
        """Detects @stencil via from-import."""
        code = """
from numba import stencil

@stencil
def my_kernel(a):
    return a[0, 0] + a[0, 1]
"""
        assert contains_jit_decorator(code)

    def test_numba_cfunc(self) -> None:
        """Detects @cfunc via from-import with arguments."""
        code = """
from numba import cfunc

@cfunc("float64(float64)")
def my_func(x):
    return x * 2
"""
        assert contains_jit_decorator(code)

    def test_numba_generated_jit(self) -> None:
        """Detects @generated_jit via from-import."""
        code = """
from numba import generated_jit

@generated_jit
def my_func(x):
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_cuda_jit(self) -> None:
        """Detects @numba.cuda.jit with chained attribute."""
        code = """
import numba

@numba.cuda.jit
def my_kernel():
    pass
"""
        assert contains_jit_decorator(code)

    def test_numba_cuda_jit_with_alias(self) -> None:
        """Detects @nb.cuda.jit when numba is aliased."""
        code = """
import numba as nb

@nb.cuda.jit
def my_kernel():
    pass
"""
        assert contains_jit_decorator(code)


class TestContainsJitDecoratorTorch:
    """Tests for torch JIT decorator detection."""

    def test_torch_compile(self) -> None:
        """Detects @torch.compile."""
        code = """
import torch

@torch.compile
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_torch_compile_with_alias(self) -> None:
        """Detects @th.compile when torch is aliased."""
        code = """
import torch as th

@th.compile
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_torch_compile_direct_import(self) -> None:
        """Detects @compile via from torch import compile."""
        code = """
from torch import compile

@compile
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_torch_compile_with_arguments(self) -> None:
        """Detects @torch.compile(mode=...) with arguments."""
        code = """
import torch

@torch.compile(mode="reduce-overhead")
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_torch_jit_script(self) -> None:
        """Detects @torch.jit.script."""
        code = """
import torch

@torch.jit.script
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_torch_jit_script_with_alias(self) -> None:
        """Detects @th.jit.script when torch is aliased."""
        code = """
import torch as th

@th.jit.script
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_torch_jit_trace(self) -> None:
        """Detects @torch.jit.trace."""
        code = """
import torch

@torch.jit.trace
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_torch_jit_imported_then_script(self) -> None:
        """Detects @jit.script via from torch import jit."""
        code = """
from torch import jit

@jit.script
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_torch_jit_imported_then_trace(self) -> None:
        """Detects @jit.trace via from torch import jit."""
        code = """
from torch import jit

@jit.trace
def my_func():
    pass
"""
        assert contains_jit_decorator(code)


class TestContainsJitDecoratorTensorFlow:
    """Tests for TensorFlow JIT decorator detection."""

    def test_tensorflow_function_with_tf_alias(self) -> None:
        """Detects @tf.function when tensorflow is aliased as tf."""
        code = """
import tensorflow as tf

@tf.function
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_tensorflow_function_full_name(self) -> None:
        """Detects @tensorflow.function with full module name."""
        code = """
import tensorflow

@tensorflow.function
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_tensorflow_function_direct_import(self) -> None:
        """Detects @function via from tensorflow import function."""
        code = """
from tensorflow import function

@function
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_tensorflow_function_with_arguments(self) -> None:
        """Detects @tf.function(jit_compile=True) with arguments."""
        code = """
import tensorflow as tf

@tf.function(jit_compile=True)
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_tf_function_direct_import_alias(self) -> None:
        """Detects @tf_func when function is aliased."""
        code = """
from tensorflow import function as tf_func

@tf_func
def my_func():
    pass
"""
        assert contains_jit_decorator(code)


class TestContainsJitDecoratorJax:
    """Tests for JAX JIT decorator detection."""

    def test_jax_jit(self) -> None:
        """Detects @jax.jit."""
        code = """
import jax

@jax.jit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_jax_jit_with_alias(self) -> None:
        """Detects @j.jit when jax is aliased."""
        code = """
import jax as j

@j.jit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_jax_jit_direct_import(self) -> None:
        """Detects @jit via from jax import jit."""
        code = """
from jax import jit

@jit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_jax_jit_direct_import_with_alias(self) -> None:
        """Detects @jax_jit when jit is aliased."""
        code = """
from jax import jit as jax_jit

@jax_jit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_jax_jit_with_arguments(self) -> None:
        """Detects @jax.jit(static_argnums=...) with arguments."""
        code = """
import jax

@jax.jit(static_argnums=(0,))
def my_func(x, y):
    pass
"""
        assert contains_jit_decorator(code)


class TestContainsJitDecoratorNegativeCases:
    """Tests that should NOT detect JIT decorators."""

    def test_no_decorators(self) -> None:
        """Returns False when there are no decorators."""
        code = """
def my_func():
    pass
"""
        assert not contains_jit_decorator(code)

    def test_other_decorator(self) -> None:
        """Returns False for non-JIT decorator like lru_cache."""
        code = """
import functools

@functools.lru_cache
def my_func():
    pass
"""
        assert not contains_jit_decorator(code)

    def test_custom_decorator(self) -> None:
        """Returns False for a user-defined decorator."""
        code = """
def my_decorator(func):
    return func

@my_decorator
def my_func():
    pass
"""
        assert not contains_jit_decorator(code)

    def test_property_decorator(self) -> None:
        """Returns False for @property."""
        code = """
class MyClass:
    @property
    def my_prop(self):
        return self._value
"""
        assert not contains_jit_decorator(code)

    def test_staticmethod_decorator(self) -> None:
        """Returns False for @staticmethod."""
        code = """
class MyClass:
    @staticmethod
    def my_func():
        pass
"""
        assert not contains_jit_decorator(code)

    def test_classmethod_decorator(self) -> None:
        """Returns False for @classmethod."""
        code = """
class MyClass:
    @classmethod
    def my_func(cls):
        pass
"""
        assert not contains_jit_decorator(code)

    def test_jit_in_comment(self) -> None:
        """Returns False when JIT decorator is in a comment."""
        code = """
# @numba.jit
def my_func():
    pass
"""
        assert not contains_jit_decorator(code)

    def test_jit_in_string(self) -> None:
        """Returns False when JIT decorator is in a docstring."""
        code = '''
def my_func():
    """This function could use @numba.jit decorator."""
    pass
'''
        assert not contains_jit_decorator(code)

    def test_unrelated_jit_name(self) -> None:
        """Returns False for locally-defined function named jit."""
        code = """
def jit():
    pass

@jit
def my_func():
    pass
"""
        assert not contains_jit_decorator(code)

    def test_unrelated_module_with_jit_attribute(self) -> None:
        """Returns False for @my_module.jit from an unrelated module."""
        code = """
import my_module

@my_module.jit
def my_func():
    pass
"""
        assert not contains_jit_decorator(code)

    def test_numba_import_but_no_decorator(self) -> None:
        """Returns False when numba is imported but not used as decorator."""
        code = """
import numba

def my_func():
    pass
"""
        assert not contains_jit_decorator(code)

    def test_jit_variable_not_decorator(self) -> None:
        """Returns False when jit is used as a variable, not decorator."""
        code = """
from numba import jit

def my_func():
    x = jit
    pass
"""
        assert not contains_jit_decorator(code)


class TestContainsJitDecoratorEdgeCases:
    """Edge case tests for JIT decorator detection."""

    def test_multiple_decorators_with_jit(self) -> None:
        """Detects JIT when stacked with other decorators."""
        code = """
import numba
import functools

@functools.lru_cache
@numba.jit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_multiple_decorators_jit_first(self) -> None:
        """Detects JIT when it is the outermost decorator."""
        code = """
import numba
import functools

@numba.jit
@functools.lru_cache
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_async_function_with_jit(self) -> None:
        """Returns False for async function with JIT decorator."""
        code = """
import numba

@numba.jit
async def my_func():
    pass
"""
        assert contains_jit_decorator(code) is False

    def test_method_in_class_with_jit(self) -> None:
        """Detects JIT on a class method."""
        code = """
import numba

class MyClass:
    @numba.jit
    def my_method(self):
        pass
"""
        assert contains_jit_decorator(code)

    def test_nested_class_method_with_jit(self) -> None:
        """Detects JIT on a nested class method."""
        code = """
import numba

class Outer:
    class Inner:
        @numba.jit
        def my_method(self):
            pass
"""
        assert contains_jit_decorator(code)

    def test_multiple_functions_one_with_jit(self) -> None:
        """Detects JIT when only one function is decorated."""
        code = """
import numba

def func_a():
    pass

@numba.jit
def func_b():
    pass

def func_c():
    pass
"""
        assert contains_jit_decorator(code)

    def test_multiple_jit_functions(self) -> None:
        """Detects JIT when multiple functions have JIT decorators."""
        code = """
import numba
import jax

@numba.jit
def func_a():
    pass

@jax.jit
def func_b():
    pass
"""
        assert contains_jit_decorator(code)

    def test_empty_code(self) -> None:
        """Returns False for empty string."""
        code = ""
        assert not contains_jit_decorator(code)

    def test_syntax_error_code(self) -> None:
        """Returns False for code with syntax errors."""
        code = """
def func(
    pass
"""
        assert not contains_jit_decorator(code)

    def test_whitespace_only(self) -> None:
        """Returns False for whitespace-only input."""
        code = "   \n\n   \t\t\n"
        assert not contains_jit_decorator(code)

    def test_only_imports(self) -> None:
        """Returns False when only imports are present."""
        code = """
import numba
from jax import jit
"""
        assert not contains_jit_decorator(code)

    def test_lambda_cannot_have_decorator(self) -> None:
        """Returns False for lambda (lambdas cannot have decorators)."""
        code = """
import numba

f = lambda x: x * 2
"""
        assert not contains_jit_decorator(code)

    def test_mixed_imports_and_aliases(self) -> None:
        """Detects JIT with mixed import styles and aliases."""
        code = """
import numba as nb
from torch import compile as torch_compile
import jax

@nb.jit
def func_a():
    pass
"""
        assert contains_jit_decorator(code)

    def test_decorator_in_different_module_context(self) -> None:
        """Detects JIT in a class method with surrounding code."""
        code = """
# Import numba for numeric computation
import numba

# Some other code
x = 5

class Processor:
    @numba.njit
    def process(self, data):
        return data * 2
"""
        assert contains_jit_decorator(code)

    def test_from_import_star_not_tracked(self) -> None:
        """Returns False for star import (not tracked)."""
        code = """
from numba import *

@jit
def my_func():
    pass
"""
        # Star imports are not tracked, so this returns False
        assert not contains_jit_decorator(code)

    def test_multiple_from_imports_same_module(self) -> None:
        """Detects JIT with multiple from-imports from same module."""
        code = """
from numba import jit
from numba import njit

@njit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)

    def test_reimport_with_different_alias(self) -> None:
        """Detects JIT when same symbol is reimported with a new alias."""
        code = """
from numba import jit
from numba import jit as fast_jit

@fast_jit
def my_func():
    pass
"""
        assert contains_jit_decorator(code)


class TestContainsJitDecoratorComplexCases:
    """Complex real-world scenarios for JIT decorator detection."""

    def test_realistic_numba_code(self) -> None:
        """Detects JIT in realistic numba code with prange."""
        code = """
import numpy as np
from numba import jit, prange

@jit(nopython=True, parallel=True)
def compute_sum(arr):
    total = 0.0
    for i in prange(len(arr)):
        total += arr[i]
    return total

def main():
    data = np.random.rand(1000000)
    result = compute_sum(data)
    print(result)
"""
        assert contains_jit_decorator(code)

    def test_realistic_torch_code(self) -> None:
        """Detects JIT in realistic torch model with @torch.compile."""
        code = """
import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)

    @torch.compile
    def forward(self, x):
        return self.linear(x)
"""
        assert contains_jit_decorator(code)

    def test_realistic_jax_code(self) -> None:
        """Detects JIT in realistic JAX loss function with grad."""
        code = """
import jax
import jax.numpy as jnp
from jax import jit, grad

@jit
def loss_fn(params, x, y):
    pred = jnp.dot(x, params)
    return jnp.mean((pred - y) ** 2)

grad_fn = grad(loss_fn)
"""
        assert contains_jit_decorator(code)

    def test_realistic_tensorflow_code(self) -> None:
        """Detects JIT in realistic TensorFlow training step."""
        code = """
import tensorflow as tf

@tf.function(jit_compile=True)
def train_step(model, x, y):
    with tf.GradientTape() as tape:
        predictions = model(x)
        loss = tf.reduce_mean(tf.square(predictions - y))
    gradients = tape.gradient(loss, model.trainable_variables)
    return loss, gradients
"""
        assert contains_jit_decorator(code)

    def test_file_with_many_functions_one_jit(self) -> None:
        """Detects JIT when one function among many is decorated."""
        code = """
import os
import sys
import numpy as np
from numba import njit

def helper_a():
    return 1

def helper_b():
    return 2

class DataProcessor:
    def __init__(self):
        self.data = []

    def process(self):
        pass

@njit
def fast_compute(x, y):
    return x + y

def main():
    result = fast_compute(1, 2)
    print(result)

if __name__ == "__main__":
    main()
"""
        assert contains_jit_decorator(code)
