"""Tests for lazy toggle-gated decoration architecture.

Covers: lazy wrapping lifecycle, toggle state, detached imports, sys.modules
crawl, permanent model preparation, pause_logging, functools.wraps
transparency, decorated identity, JIT compatibility, wrap/unwrap, and
exception/interrupt safety.
"""

import gc
import inspect
import os
import signal
import sys
import types
import weakref

import pytest
import torch
from torch import nn

import torchlens
from torchlens.decoration import torch_funcs as torch_funcs_module
from torchlens import _state, log_forward_pass
from torchlens.decoration.torch_funcs import (
    decorate_all_once,
    get_func_argnames,
    patch_detached_references,
    patch_model_instance,
    wrap_torch,
    unwrap_torch,
    wrapped,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


class SimpleModel(nn.Module):
    def forward(self, x):
        return torch.relu(x + 1)


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

    def forward(self, x):
        return torch.sigmoid(self.linear(x))


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

    def forward(self, x, y):
        return self.linear(x + y)


class ErrorModel(nn.Module):
    def forward(self, x):
        torch.relu(x)
        raise RuntimeError("intentional error")


class KeyboardInterruptModel(nn.Module):
    def forward(self, x):
        torch.relu(x)
        raise KeyboardInterrupt("simulated ctrl-c")


class InplaceModel(nn.Module):
    def forward(self, x):
        x = x.clone()
        x.add_(1)
        x.relu_()
        return x


class PropertyAccessModel(nn.Module):
    def forward(self, x):
        return x.real + x.T.sum()


class InternalCreationModel(nn.Module):
    def forward(self, x):
        y = torch.zeros(3)
        return x.sum() + y.sum()


class NestedModuleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.block = nn.Sequential(nn.Linear(5, 5), nn.ReLU())

    def forward(self, x):
        return self.block(x)


class IdentityModuleModel(nn.Module):
    """Model with nn.Identity — tests decorated identity injection."""

    def __init__(self):
        super().__init__()
        self.id = nn.Identity()
        self.linear = nn.Linear(5, 5)

    def forward(self, x):
        return self.linear(self.id(x))


# =========================================================================
# 0. Lazy Decoration Lifecycle
# =========================================================================


class TestLazyDecoration:
    @pytest.fixture(autouse=True)
    def _clean_state(self):
        """Unwrap before each test, re-wrap after for other test classes."""
        unwrap_torch()
        yield
        wrap_torch()

    @pytest.mark.smoke
    def test_import_does_not_decorate(self):
        """After import torchlens (with unwrapped state), torch functions are originals."""
        assert _state._is_decorated is False
        assert not getattr(torch.cos, "tl_is_decorated_function", False)

    @pytest.mark.smoke
    def test_log_forward_pass_triggers_wrapping(self):
        """First log_forward_pass should auto-wrap torch functions."""
        assert _state._is_decorated is False
        model = SimpleModel()
        log_forward_pass(model, torch.randn(5))
        assert _state._is_decorated is True
        assert getattr(torch.cos, "tl_is_decorated_function", False)

    def test_stays_wrapped_after_log_forward_pass(self):
        """After log_forward_pass, torch stays wrapped for subsequent calls."""
        model = SimpleModel()
        log_forward_pass(model, torch.randn(5))
        assert _state._is_decorated is True

        # Second call should also work without issues
        log_forward_pass(model, torch.randn(5))
        assert _state._is_decorated is True

    def test_unwrap_when_done_parameter(self):
        """log_forward_pass(unwrap_when_done=True) restores originals after."""
        model = SimpleModel()
        result = log_forward_pass(model, torch.randn(5), unwrap_when_done=True)
        assert result is not None
        assert len(result.output_layers) > 0
        assert _state._is_decorated is False
        assert not getattr(torch.cos, "tl_is_decorated_function", False)

    def test_wrapped_context_manager(self):
        """wrapped() context manager wraps on entry, unwraps on exit."""
        assert _state._is_decorated is False
        with wrapped():
            assert _state._is_decorated is True
            model = SimpleModel()
            result = log_forward_pass(model, torch.randn(5))
            assert len(result.output_layers) > 0
        assert _state._is_decorated is False

    def test_wrapped_context_manager_exception_safety(self):
        """wrapped() must unwrap even if body raises."""
        assert _state._is_decorated is False
        with pytest.raises(RuntimeError):
            with wrapped():
                assert _state._is_decorated is True
                raise RuntimeError("test error")
        assert _state._is_decorated is False

    def test_wrap_unwrap_cycle(self):
        """Multiple wrap/unwrap cycles should work correctly."""
        for _ in range(3):
            wrap_torch()
            assert _state._is_decorated is True
            assert getattr(torch.cos, "tl_is_decorated_function", False)
            unwrap_torch()
            assert _state._is_decorated is False
            assert not getattr(torch.cos, "tl_is_decorated_function", False)

    def test_log_forward_pass_after_unwrap(self):
        """log_forward_pass after explicit unwrap should auto-rewrap and work."""
        wrap_torch()
        unwrap_torch()
        assert _state._is_decorated is False

        model = SimpleModel()
        result = log_forward_pass(model, torch.randn(5))
        assert _state._is_decorated is True
        assert len(result.output_layers) > 0

    def test_unwrap_is_idempotent(self):
        """Calling unwrap_torch multiple times should not crash."""
        unwrap_torch()
        unwrap_torch()
        assert _state._is_decorated is False

    def test_wrap_is_idempotent(self):
        """Calling wrap_torch multiple times should not crash."""
        wrap_torch()
        wrap_torch()
        assert _state._is_decorated is True


# =========================================================================
# 1. Toggle State
# =========================================================================


class TestToggleState:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    @pytest.mark.smoke
    def test_toggle_off_by_default(self):
        """After import, logging toggle should be off."""
        assert _state._logging_enabled is False
        assert _state._active_model_log is None

    @pytest.mark.smoke
    def test_toggle_on_during_logging(self):
        """Inside log_forward_pass, toggle should be True."""
        observed = {}

        class ProbeModel(nn.Module):
            def forward(self, x):
                observed["enabled"] = _state._logging_enabled
                observed["model_log"] = _state._active_model_log
                return x + 1

        model = ProbeModel()
        log_forward_pass(model, torch.randn(3))
        assert observed["enabled"] is True
        assert observed["model_log"] is not None

    @pytest.mark.smoke
    def test_toggle_off_after_logging(self):
        """After log_forward_pass completes, toggle should be off."""
        model = SimpleModel()
        log_forward_pass(model, torch.randn(5))
        assert _state._logging_enabled is False
        assert _state._active_model_log is None

    def test_toggle_off_after_exception(self):
        """After an exception in forward, toggle must be off."""
        model = ErrorModel()
        with pytest.raises(RuntimeError, match="intentional error"):
            log_forward_pass(model, torch.randn(5))
        assert _state._logging_enabled is False
        assert _state._active_model_log is None

    def test_toggle_off_after_keyboard_interrupt(self):
        """After KeyboardInterrupt in forward, toggle must be off."""
        model = KeyboardInterruptModel()
        with pytest.raises(KeyboardInterrupt):
            log_forward_pass(model, torch.randn(5))
        assert _state._logging_enabled is False
        assert _state._active_model_log is None

    def test_torch_works_after_keyboard_interrupt(self):
        """Torch functions must work normally after a KeyboardInterrupt."""
        model = KeyboardInterruptModel()
        with pytest.raises(KeyboardInterrupt):
            log_forward_pass(model, torch.randn(5))
        # These must not crash or produce tl_ attributes
        x = torch.randn(3)
        y = torch.cos(x)
        assert y.shape == (3,)
        assert not hasattr(y, "tl_tensor_label_raw")

    def test_requires_grad_restored_after_exception(self):
        """requires_grad must be restored even when forward raises."""
        model = TwoLayerModel()
        orig_grads = {n: p.requires_grad for n, p in model.named_parameters()}
        with pytest.raises(RuntimeError):
            # Force an error by passing wrong shape
            log_forward_pass(model, torch.randn(0))
        for name, param in model.named_parameters():
            assert param.requires_grad == orig_grads[name], f"{name} requires_grad changed"


class TestWrapUnwrap:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Start wrapped, ensure wrapped on exit."""
        wrap_torch()
        yield
        wrap_torch()

    def test_unwrap_restores_original_torch_function(self):
        """unwrap_torch strips TorchLens wrappers from torch callables."""
        assert getattr(torch.cos, "tl_is_decorated_function", False)
        unwrap_torch()
        assert not getattr(torch.cos, "tl_is_decorated_function", False)
        x = torch.randn(4)
        y = torch.cos(x)
        assert y.shape == x.shape
        assert not hasattr(y, "tl_tensor_label_raw")

    def test_wrap_restores_logging_after_unwrap(self):
        """wrap_torch after unwrap_torch re-enables logging."""
        unwrap_torch()
        assert not getattr(torch.nn.functional.relu, "tl_is_decorated_function", False)
        wrap_torch()
        assert getattr(torch.nn.functional.relu, "tl_is_decorated_function", False)
        result = log_forward_pass(SimpleModel(), torch.randn(5))
        relu_layers = [label for label in result.layer_labels if "relu" in label.lower()]
        assert relu_layers


# =========================================================================
# 1b. Sequence Slot Fix (sq_item pollution from __getitem__ wrapping)
# =========================================================================


class TestSequenceSlotFix:
    """Wrapping Tensor.__getitem__ pollutes CPython's sq_item C slot, making
    PySequence_Check(tensor) return True.  This breaks torch.tensor() on lists
    of 0-d tensors.  These tests verify the fix holds across all lifecycle paths.
    """

    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Start wrapped, ensure wrapped on exit."""
        wrap_torch()
        yield
        wrap_torch()

    @staticmethod
    def _check_tensor_from_0d(msg: str = "") -> None:
        """Assert torch.tensor([0-d, 0-d]) works without TypeError."""
        D = torch.randn(3, 3)
        result = torch.tensor([D[0, 0], D[1, 1]])
        assert result.shape == (2,), f"wrong shape {msg}"
        assert torch.allclose(result, torch.stack([D[0, 0], D[1, 1]])), msg

    @pytest.mark.smoke
    def test_tensor_from_0d_while_wrapped(self):
        """torch.tensor([0-d, 0-d]) must work while decoration is active."""
        self._check_tensor_from_0d("while wrapped")

    @pytest.mark.smoke
    def test_tensor_from_0d_after_unwrap(self):
        """torch.tensor([0-d, 0-d]) must work after unwrap_torch()."""
        unwrap_torch()
        self._check_tensor_from_0d("after unwrap")

    def test_tensor_from_0d_after_wrap_unwrap_cycle(self):
        """torch.tensor([0-d, 0-d]) survives multiple wrap/unwrap cycles."""
        for i in range(3):
            unwrap_torch()
            self._check_tensor_from_0d(f"unwrap cycle {i}")
            wrap_torch()
            self._check_tensor_from_0d(f"wrap cycle {i}")

    def test_tensor_from_0d_wrapped_context_manager(self):
        """torch.tensor([0-d, 0-d]) works inside and after wrapped()."""
        unwrap_torch()
        with wrapped():
            self._check_tensor_from_0d("inside wrapped()")
        self._check_tensor_from_0d("after wrapped() exit")

    def test_tensor_from_0d_after_forward_pass(self):
        """torch.tensor([0-d, 0-d]) works after a real forward pass."""
        model = SimpleModel()
        log_forward_pass(model, torch.randn(5))
        self._check_tensor_from_0d("after forward pass")

    def test_tensor_from_0d_after_unwrap_when_done(self):
        """torch.tensor([0-d, 0-d]) works after log_forward_pass(unwrap_when_done=True)."""
        model = SimpleModel()
        log_forward_pass(model, torch.randn(5), unwrap_when_done=True)
        self._check_tensor_from_0d("after unwrap_when_done")

    def test_tensor_from_0d_nested_list(self):
        """torch.tensor with nested lists of 0-d tensors also works."""
        D = torch.randn(2, 2)
        result = torch.tensor([[D[0, 0], D[0, 1]], [D[1, 0], D[1, 1]]])
        assert result.shape == (2, 2)
        assert torch.allclose(result, D)

    def test_tensor_indexing_still_works(self):
        """Clearing sq_item must not break normal tensor indexing."""
        x = torch.randn(3, 4, 5)
        assert x[0].shape == (4, 5)
        assert x[0, 1].shape == (5,)
        assert x[0, 1, 2].shape == ()
        assert x[:, 1:3].shape == (3, 2, 5)
        assert x[torch.tensor([0, 2])].shape == (2, 4, 5)

    def test_sequence_check_false_for_tensors(self):
        """PySequence_Check must return False for tensors after decoration."""
        import ctypes

        if sys.implementation.name != "cpython":
            pytest.skip("ctypes slot check only works on CPython")
        check = ctypes.pythonapi.PySequence_Check
        check.argtypes = [ctypes.py_object]
        check.restype = ctypes.c_int
        t = torch.tensor(0.5)
        assert check(t) == 0, "PySequence_Check(tensor) should be False"


# =========================================================================
# 2. Torch Functions Normal When Toggle Off
# =========================================================================


class TestPassthroughWhenOff:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_basic_ops_produce_correct_results(self):
        """Decorated torch functions must produce identical results when off."""
        x = torch.tensor([1.0, 2.0, 3.0])
        assert torch.allclose(torch.cos(x), x.cos())
        assert torch.allclose(torch.sin(x), x.sin())
        assert torch.allclose(torch.relu(x), torch.nn.functional.relu(x))
        assert torch.allclose(x + x, torch.add(x, x))

    def test_no_tl_attributes_on_results(self):
        """Results must not have tl_ attributes when toggle is off."""
        x = torch.randn(3, 4)
        y = torch.matmul(x, x.T)
        assert not hasattr(y, "tl_tensor_label_raw")
        assert not hasattr(y, "tl_source_model_log")

    def test_ops_normal_after_logging(self):
        """Torch ops must be clean after a logging session completes."""
        model = SimpleModel()
        log_forward_pass(model, torch.randn(5))

        x = torch.randn(5)
        y = torch.relu(x)
        z = x.view(1, 5)
        assert not hasattr(y, "tl_tensor_label_raw")
        assert not hasattr(z, "tl_tensor_label_raw")

    def test_tensor_method_ops(self):
        """Tensor method operations like .view, .reshape must work."""
        x = torch.randn(6)
        assert x.view(2, 3).shape == (2, 3)
        assert x.reshape(3, 2).shape == (3, 2)
        assert x.unsqueeze(0).shape == (1, 6)


# =========================================================================
# 3. Detached Import Patching
# =========================================================================


class TestDetachedImports:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_module_level_import_patched(self):
        """A 'from torch import cos' at module level should be patched."""
        # Create a synthetic module that simulates 'from torch import cos'
        mod = types.ModuleType("_test_detached_cos")
        mod.cos = torch.cos  # torch.cos is already decorated at this point
        sys.modules["_test_detached_cos"] = mod
        try:
            # The function should be decorated
            assert getattr(mod.cos, "tl_is_decorated_function", False)
        finally:
            del sys.modules["_test_detached_cos"]

    def test_model_with_stored_torch_func(self):
        """A model storing self.act = torch.relu should have it patched."""

        class FuncAttrModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.act = torch.relu
                self.linear = nn.Linear(5, 5)

            def forward(self, x):
                return self.act(self.linear(x))

        model = FuncAttrModel()
        # Before patching, self.act might be undecorated (if bound to original)
        # patch_model_instance should fix it
        patch_model_instance(model)
        # After patching, the relu stored on the model should be decorated
        result = log_forward_pass(model, torch.randn(5))
        # The relu should appear in the graph
        relu_layers = [label for label in result.layer_labels if "relu" in label.lower()]
        assert len(relu_layers) > 0, "relu from self.act not logged in graph"

    def test_model_with_func_in_list(self):
        """A model storing functions in a list should still have them logged."""

        class ListFuncModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.funcs = [torch.relu, torch.sigmoid]
                self.linear = nn.Linear(5, 5)

            def forward(self, x):
                x = self.linear(x)
                for f in self.funcs:
                    x = f(x)
                return x

        model = ListFuncModel()
        result = log_forward_pass(model, torch.randn(5))
        labels = " ".join(result.layer_labels).lower()
        assert "relu" in labels, "relu from list not logged"
        assert "sigmoid" in labels, "sigmoid from list not logged"

    def test_model_with_func_in_dict(self):
        """A model storing functions in a dict should still have them logged."""

        class DictFuncModel(nn.Module):
            def __init__(self):
                super().__init__()
                self.ops = {"activation": torch.relu}
                self.linear = nn.Linear(5, 5)

            def forward(self, x):
                return self.ops["activation"](self.linear(x))

        model = DictFuncModel()
        result = log_forward_pass(model, torch.randn(5))
        relu_layers = [lbl for lbl in result.layer_labels if "relu" in lbl.lower()]
        assert len(relu_layers) > 0

    def test_nn_functional_import_patched(self):
        """torch.nn.functional functions should be decorated."""
        import torch.nn.functional as F

        assert getattr(F.relu, "tl_is_decorated_function", False)
        assert getattr(F.linear, "tl_is_decorated_function", False)

    def test_late_import_patched_incrementally(self):
        """Modules imported after torchlens should be patched on next crawl."""
        mod_name = "_test_late_import_module"
        # Remove if somehow already present
        sys.modules.pop(mod_name, None)
        _state._crawled_module_keys.discard(mod_name)

        # Simulate a late import
        mod = types.ModuleType(mod_name)
        # Store the DECORATED cos (since torch.cos is already decorated)
        mod.my_cos = torch.cos
        sys.modules[mod_name] = mod
        try:
            # Trigger incremental crawl
            patch_detached_references()
            assert getattr(mod.my_cos, "tl_is_decorated_function", False)
        finally:
            sys.modules.pop(mod_name, None)

    def test_crawl_skips_torchlens_modules(self):
        """The crawl must not modify torchlens internal modules."""
        # torchlens modules should be in _crawled_module_keys but NOT patched
        tl_modules = [k for k in sys.modules if k.startswith("torchlens")]
        assert len(tl_modules) > 0
        # _state itself should not have been modified by the crawl
        assert not hasattr(_state, "tl_is_decorated_function")

    def test_crawl_only_processes_new_modules(self):
        """Calling patch_detached_references twice should not re-scan."""
        keys_after_first = set(_state._crawled_module_keys)
        patch_detached_references()
        keys_after_second = set(_state._crawled_module_keys)
        # If no new modules were imported, sets should be identical
        assert keys_after_first == keys_after_second


# =========================================================================
# 4. Permanent Model Preparation
# =========================================================================


class TestPermanentModelPrep:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_model_prepared_once(self):
        """Same model instance should only be prepared once."""
        model = TwoLayerModel()
        log_forward_pass(model, torch.randn(5))
        assert model in _state._prepared_models

        # Run again — should reuse
        log_forward_pass(model, torch.randn(5))
        assert model in _state._prepared_models

    def test_permanent_attrs_survive_sessions(self):
        """tl_module_address and tl_module_type should persist after logging."""
        model = NestedModuleModel()
        log_forward_pass(model, torch.randn(5))

        # Permanent attributes should still be there
        assert hasattr(model.block[0], "tl_module_address")
        assert hasattr(model.block[0], "tl_module_type")
        assert model.block[0].tl_module_address == "block.0"
        assert model.block[0].tl_module_type == "Linear"

    def test_session_attrs_cleaned(self):
        """Session-scoped tl_ attrs should be removed after logging."""
        model = TwoLayerModel()
        log_forward_pass(model, torch.randn(5))

        for module in model.modules():
            assert not hasattr(module, "tl_source_model_log")
            assert not hasattr(module, "tl_module_pass_num")
            assert not hasattr(module, "tl_tensors_entered_labels")
            assert not hasattr(module, "tl_tensors_exited_labels")

    def test_session_attrs_cleaned_after_exception(self):
        """Session attrs must be cleaned even if forward raises."""
        model = nn.Sequential(nn.Linear(5, 5), nn.ReLU())

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

            def forward(self, x):
                x = self.linear(x)
                raise RuntimeError("boom")

        model = FailAfterLinear()
        with pytest.raises(RuntimeError):
            log_forward_pass(model, torch.randn(5))

        for module in model.modules():
            assert not hasattr(module, "tl_source_model_log")

    def test_different_models_prepared_independently(self):
        """Two different model instances should both be in _prepared_models."""
        model_a = SimpleModel()
        model_b = TwoLayerModel()
        log_forward_pass(model_a, torch.randn(5))
        log_forward_pass(model_b, torch.randn(5))
        assert model_a in _state._prepared_models
        assert model_b in _state._prepared_models

    def test_model_gc_after_deletion(self):
        """WeakSet should not prevent model garbage collection.

        Note: log_forward_pass returns a ModelLog that holds refs to the model
        (known GC issues GC-1 through GC-5). We test that the WeakSet itself
        doesn't add a strong ref by preparing a model WITHOUT logging.
        """
        from torchlens.decoration.model_prep import _prepare_model_once

        model = TwoLayerModel()
        ref = weakref.ref(model)
        _prepare_model_once(model)
        assert model in _state._prepared_models

        del model
        gc.collect()
        assert ref() is None  # model should be collected

    def test_requires_grad_restored(self):
        """All params should have original requires_grad after logging."""
        model = TwoLayerModel()
        # Freeze some params
        model.linear.bias.requires_grad_(False)
        orig_grads = {n: p.requires_grad for n, p in model.named_parameters()}

        log_forward_pass(model, torch.randn(5))

        for name, param in model.named_parameters():
            assert param.requires_grad == orig_grads[name], (
                f"{name}: expected requires_grad={orig_grads[name]}"
            )


# =========================================================================
# 5. pause_logging Context Manager
# =========================================================================


class TestPauseLogging:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_pause_restores_state(self):
        """pause_logging should restore previous _logging_enabled value."""
        _state._logging_enabled = True
        try:
            with _state.pause_logging():
                assert _state._logging_enabled is False
            assert _state._logging_enabled is True
        finally:
            _state._logging_enabled = False

    def test_pause_nested(self):
        """Nested pause_logging should correctly restore at each level."""
        _state._logging_enabled = True
        try:
            with _state.pause_logging():
                assert _state._logging_enabled is False
                with _state.pause_logging():
                    assert _state._logging_enabled is False
                assert _state._logging_enabled is False  # restored to inner's prev
            assert _state._logging_enabled is True  # restored to outer's prev
        finally:
            _state._logging_enabled = False

    def test_pause_exception_safety(self):
        """pause_logging must restore state even if body raises."""
        _state._logging_enabled = True
        try:
            with pytest.raises(ValueError):
                with _state.pause_logging():
                    assert _state._logging_enabled is False
                    raise ValueError("test")
            assert _state._logging_enabled is True
        finally:
            _state._logging_enabled = False

    def test_operations_during_pause_not_logged(self):
        """Torch ops inside pause_logging should not appear in graph."""

        class PauseInsideModel(nn.Module):
            def forward(self, x):
                y = torch.relu(x)  # should be logged
                with _state.pause_logging():
                    torch.cos(x)  # should NOT be logged
                return y

        model = PauseInsideModel()
        result = log_forward_pass(model, torch.randn(5))
        labels = " ".join(result.layer_labels).lower()
        assert "relu" in labels
        # cos should not appear — it ran during pause
        cos_layers = [lbl for lbl in result.layer_labels if "cos" in lbl.lower()]
        assert len(cos_layers) == 0, f"cos appeared in graph during pause: {cos_layers}"

    def test_activation_postfunc_not_logged(self):
        """activation_postfunc runs inside pause_logging, ops should not appear."""
        model = SimpleModel()
        result = log_forward_pass(model, torch.randn(5), activation_postfunc=torch.sigmoid)
        # sigmoid from postfunc should NOT appear in graph
        sigmoid_layers = [lbl for lbl in result.layer_labels if "sigmoid" in lbl.lower()]
        assert len(sigmoid_layers) == 0, f"postfunc sigmoid leaked into graph: {sigmoid_layers}"


# =========================================================================
# 6. functools.wraps / Transparency
# =========================================================================


class TestWrapperTransparency:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_wrapped_function_name(self):
        """Decorated functions should preserve __name__."""
        assert torch.cos.__name__ == "cos"
        assert torch.nn.functional.relu.__name__ == "relu"

    def test_wrapped_function_doc(self):
        """Decorated functions should preserve __doc__."""
        assert torch.cos.__doc__ is not None
        assert len(torch.cos.__doc__) > 0

    def test_is_decorated_flag(self):
        """Decorated functions should have tl_is_decorated_function=True."""
        assert getattr(torch.cos, "tl_is_decorated_function", False)
        assert getattr(torch.nn.functional.relu, "tl_is_decorated_function", False)

    def test_builtin_no_wrapped(self):
        """Built-in wrappers should NOT have __wrapped__ (JIT compat)."""
        # Built-ins like torch.cos don't have __code__; we remove __wrapped__
        # to prevent inspect.unwrap from following it.
        if not hasattr(torch.cos, "__code__"):
            assert not hasattr(torch.cos, "__wrapped__")


# =========================================================================
# 7. Decorated Identity (via _state._decorated_identity)
# =========================================================================


class TestDecoratedIdentity:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_decorated_identity_exists(self):
        """_state._decorated_identity should be set after wrapping."""
        assert _state._decorated_identity is not None
        assert callable(_state._decorated_identity)

    def test_decorated_identity_passthrough(self):
        """_decorated_identity should return the input unchanged."""
        x = torch.randn(3, 4)
        y = _state._decorated_identity(x)
        assert torch.equal(x, y)

    def test_identity_not_on_torch(self):
        """torch.identity should NOT exist — we don't monkey-patch torch."""
        assert not hasattr(torch, "identity")

    def test_identity_in_graph(self):
        """nn.Identity modules should produce identity entries in the graph."""
        model = IdentityModuleModel()
        result = log_forward_pass(model, torch.randn(5))
        identity_layers = [lbl for lbl in result.layer_labels if "identity" in lbl.lower()]
        assert len(identity_layers) > 0

    def test_decorated_identity_cleared_on_unwrap(self):
        """After unwrap_torch, _decorated_identity should still be callable
        (the wrapper object persists) but _is_decorated should be False."""
        unwrap_torch()
        assert _state._is_decorated is False
        # Re-wrap for other tests
        wrap_torch()


# =========================================================================
# 8. JIT Compatibility
# =========================================================================


class TestJITCompat:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_jit_builtin_table_has_decorated(self):
        """Decorated functions should be in the JIT builtin table."""
        import torch.jit._builtins as _jit_builtins

        # torch.cos should be registered
        assert id(torch.cos) in _jit_builtins._builtin_table

    def test_shared_originals_same_wrapper(self):
        """Functions sharing an original (e.g. torch.cos / torch._VF.cos)
        should use the exact same wrapper object."""
        if hasattr(torch._VF, "cos"):
            assert torch.cos is torch._VF.cos

    def test_jit_script_with_decorated_functions(self):
        """torch.jit.script should work on functions using decorated ops."""

        @torch.jit.script
        def jit_fn(x: torch.Tensor) -> torch.Tensor:
            return torch.relu(x) + 1.0

        x = torch.randn(5)
        result = jit_fn(x)
        expected = torch.relu(x) + 1.0
        assert torch.allclose(result, expected)


# =========================================================================
# 9. Decoration Mapping Consistency
# =========================================================================


class TestFuncArgnames:
    def test_get_func_argnames_tolerates_bad_annotations(
        self, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        """get_func_argnames should tolerate ``inspect.signature`` TypeErrors. #138"""

        def fake_func(x: object, *, flag: bool = False) -> None:
            """fake_func(x, *, flag=False)"""

        original_signature = torch_funcs_module.inspect.signature

        def failing_signature(target: object) -> inspect.Signature:
            """Raise the simulated PEP 649 failure for the fake function only."""
            if target is fake_func:
                raise TypeError("simulated PEP 649 failure")
            return original_signature(target)

        monkeypatch.setattr(torch_funcs_module.inspect, "signature", failing_signature)
        _state._func_argnames.pop("fake_func", None)

        get_func_argnames(fake_func, "fake_func")

        assert _state._func_argnames["fake_func"] == ("x", "flag")

    def test_argnames_collected_before_decoration(self) -> None:
        """Argnames must be collected before decoration mutates torch names. #138"""
        wrap_torch()

        assert len(_state._func_argnames) > 500
        assert len(_state._orig_to_decorated) > 1000


class TestPartialDecorationRecovery:
    """Guard against partial decoration leaving the system in an inconsistent
    state.  The idempotency check in ``decorate_all_once`` must use
    ``_is_decorated`` (completion flag), NOT ``_orig_to_decorated`` (non-empty
    dict), so that a midway failure allows retry. #138"""

    def test_is_decorated_false_before_wrapping(self) -> None:
        """Before any wrapping, _is_decorated must be False."""
        unwrap_torch()
        assert _state._is_decorated is False

    def test_is_decorated_true_after_wrapping(self) -> None:
        """After successful wrapping, _is_decorated must be True."""
        wrap_torch()
        assert _state._is_decorated is True

    def test_is_decorated_distinguishes_partial_from_complete(self) -> None:
        """_is_decorated must be False when _orig_to_decorated is non-empty
        but decoration was not completed (simulated partial failure)."""
        wrap_torch()
        # Simulate partial state: decorated flag cleared, but maps remain
        _state._is_decorated = False
        assert len(_state._orig_to_decorated) > 0
        assert _state._is_decorated is False
        # decorate_all_once should NOT early-return on non-empty dict
        decorate_all_once()
        assert _state._is_decorated is True
        # Restore clean state
        wrap_torch()

    def test_unwrap_clears_is_decorated_but_keeps_maps(self) -> None:
        """unwrap_torch sets _is_decorated=False but preserves maps for
        re-installation.  This is the legitimate non-empty-dict case."""
        wrap_torch()
        assert _state._is_decorated is True
        unwrap_torch()
        assert _state._is_decorated is False
        assert len(_state._orig_to_decorated) > 0
        assert len(_state._decorated_to_orig) > 0
        # Re-wrap should succeed via the re-install path
        wrap_torch()
        assert _state._is_decorated is True


class TestDecorationConsistency:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_orig_to_decorated_populated(self):
        """_orig_to_decorated should be populated after wrapping."""
        assert len(_state._orig_to_decorated) > 1000

    def test_decorated_to_orig_populated(self):
        """_decorated_to_orig should mirror _orig_to_decorated."""
        assert len(_state._decorated_to_orig) == len(_state._orig_to_decorated)

    def test_bidirectional_mapper(self):
        """_decorated_func_mapper should have dec->orig for every wrapper."""
        for orig_id, dec in _state._orig_to_decorated.items():
            if isinstance(dec, property):
                continue
            orig = _state._decorated_to_orig.get(id(dec))
            if orig is None:
                continue
            # dec -> orig must always be correct
            assert dec in _state._decorated_func_mapper
            assert _state._decorated_func_mapper[dec] is orig
            # orig -> dec may point to a different wrapper if the original
            # is shared across function names (e.g. spmm/dsmm are aliases),
            # so we only check it exists, not which wrapper it points to.
            assert orig in _state._decorated_func_mapper

    def test_no_duplicate_wrappers(self):
        """Each decorated func in _orig_to_decorated should be unique."""
        dec_ids = set()
        for dec in _state._orig_to_decorated.values():
            d_id = id(dec)
            assert d_id not in dec_ids, f"Duplicate wrapper found: {dec}"
            dec_ids.add(d_id)

    def test_func_argnames_populated(self):
        """_func_argnames should be pre-computed for many functions."""
        assert len(_state._func_argnames) > 500

    def test_decorate_all_once_idempotent(self):
        """Calling decorate_all_once again should be a no-op."""
        count_before = len(_state._orig_to_decorated)
        decorate_all_once()
        assert len(_state._orig_to_decorated) == count_before


# =========================================================================
# 10. In-Place Ops
# =========================================================================


class TestInPlaceOps:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_inplace_ops_logged(self):
        """In-place operations should appear in the logged graph."""
        model = InplaceModel()
        result = log_forward_pass(model, torch.randn(5))
        labels = " ".join(result.layer_labels).lower()
        assert "add" in labels
        assert "relu" in labels

    def test_inplace_ops_correct_output(self):
        """In-place model should produce correct numerical output."""
        model = InplaceModel()
        x = torch.randn(5)
        result = log_forward_pass(model, x)
        output = result.output_layers
        assert len(output) > 0


# =========================================================================
# 11. Property Descriptors
# =========================================================================


class TestPropertyDescriptors:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_tensor_real_access(self):
        """Accessing .real should work and be logged."""
        x = torch.randn(3)
        result = x.real
        assert torch.equal(result, x)

    def test_tensor_T_access(self):
        """Accessing .T should work."""
        x = torch.randn(3, 4)
        assert x.T.shape == (4, 3)


# =========================================================================
# 12. Edge Cases
# =========================================================================


class TestEdgeCases:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_model_with_no_torch_ops(self):
        """A model that just returns input should not crash."""

        class PassthroughModel(nn.Module):
            def forward(self, x):
                return x

        model = PassthroughModel()
        result = log_forward_pass(model, torch.randn(5))
        assert result is not None

    def test_model_with_internal_tensor_creation(self):
        """torch.zeros inside forward should be logged."""
        model = InternalCreationModel()
        result = log_forward_pass(model, torch.randn(3))
        labels = " ".join(result.layer_labels).lower()
        assert "zeros" in labels or "sum" in labels

    def test_multiple_inputs(self):
        """Models with multiple tensor inputs should work."""
        model = MultiInputModel()
        result = log_forward_pass(model, [torch.randn(5), torch.randn(5)])
        assert len(result.output_layers) > 0

    def test_nested_modules(self):
        """Nested module structure should be correctly logged."""
        model = NestedModuleModel()
        result = log_forward_pass(model, torch.randn(5))
        labels = " ".join(result.layer_labels).lower()
        assert "linear" in labels
        assert "relu" in labels

    def test_fast_mode_with_toggle(self):
        """Fast mode (layers_to_save subset) should work with permanent decoration."""
        model = TwoLayerModel()
        result = log_forward_pass(model, torch.randn(5))
        # Get a specific layer to save
        all_labels = result.layer_labels
        if len(all_labels) > 2:
            result2 = log_forward_pass(model, torch.randn(5), layers_to_save=[all_labels[1]])
            assert result2 is not None

    def test_consecutive_logging_sessions(self):
        """Multiple consecutive log_forward_pass calls should work cleanly."""
        model = SimpleModel()
        for _ in range(3):
            result = log_forward_pass(model, torch.randn(5))
            assert _state._logging_enabled is False
            assert _state._active_model_log is None
            assert len(result.output_layers) > 0


# =========================================================================
# 13. SIGINT / Signal Safety
# =========================================================================


class TestSignalSafety:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_sigalrm_during_forward(self):
        """SIGALRM during a forward pass must not leave the logging toggle stuck on."""
        if not hasattr(signal, "SIGALRM"):
            pytest.skip("SIGALRM not available on this platform")

        class AlarmModel(nn.Module):
            def forward(self, x):
                x = x + 1.0  # at least one logged op before the signal
                os.kill(os.getpid(), signal.SIGALRM)
                return x + 1.0

        def alarm_handler(signum, frame):
            raise TimeoutError("alarm fired")

        old_handler = signal.signal(signal.SIGALRM, alarm_handler)
        try:
            model = AlarmModel()
            with pytest.raises(TimeoutError):
                log_forward_pass(model, torch.randn(5))

            # Toggle MUST be off
            assert _state._logging_enabled is False
            assert _state._active_model_log is None

            # Torch must still work
            y = torch.cos(torch.tensor([1.0]))
            assert y.shape == (1,)
        finally:
            signal.signal(signal.SIGALRM, old_handler)

    def test_active_logging_context_manager_exception_safety(self):
        """active_logging must restore state on any exception type."""

        # Simulate the context manager directly
        class CustomError(BaseException):
            pass

        assert _state._logging_enabled is False
        try:
            with _state.active_logging(object()):
                assert _state._logging_enabled is True
                raise CustomError("unusual error")
        except CustomError:
            pass
        assert _state._logging_enabled is False
        assert _state._active_model_log is None


# =========================================================================
# 14. Multiple Session Isolation
# =========================================================================


class TestSessionIsolation:
    @pytest.fixture(autouse=True)
    def _ensure_wrapped(self):
        """Ensure torch functions are wrapped for these tests."""
        wrap_torch()

    def test_no_cross_session_leakage(self):
        """Tensors from one session should not leak tl_ attrs into the next."""
        model = TwoLayerModel()
        x = torch.randn(5)

        result1 = log_forward_pass(model, x)
        # After first session, input tensor should be clean
        assert not hasattr(x, "tl_tensor_label_raw")

        result2 = log_forward_pass(model, x)
        assert not hasattr(x, "tl_tensor_label_raw")

        # Both sessions should have produced valid results
        assert len(result1.output_layers) > 0
        assert len(result2.output_layers) > 0

    def test_model_log_not_retained_on_state(self):
        """_active_model_log should be None between sessions."""
        model = SimpleModel()
        log_forward_pass(model, torch.randn(5))
        assert _state._active_model_log is None

        log_forward_pass(model, torch.randn(5))
        assert _state._active_model_log is None
