"""Context-stack logic with an explicit driver model; no CUDA initialization."""
import importlib.util
from pathlib import Path
import sys
from types import ModuleType, SimpleNamespace
import unittest
from unittest.mock import patch


class ContextOwnership(unittest.TestCase):
    def setUp(self):
        self.stack = []
        stack = self.stack
        class Context:
            def __init__(self):
                self.detached = False
                stack.append(self)
            def push(self):
                if self.detached:
                    raise RuntimeError("detached")
                stack.append(self)
            @staticmethod
            def pop():
                stack.pop()
            def detach(self):
                self.detached = True
        class Device:
            @staticmethod
            def count(): return 1
            def __init__(self, device): pass
            def make_context(self): return Context()
        driver = ModuleType("pycuda.driver")
        driver.Device, driver.init = Device, lambda: None
        pycuda = ModuleType("pycuda")
        pycuda.driver = driver
        path = Path(__file__).resolve().parents[2] / "src/CelerisLab/cuda/context.py"
        spec = importlib.util.spec_from_file_location("context_under_test", path)
        self.module = importlib.util.module_from_spec(spec)
        with patch.dict(sys.modules, {"pycuda": pycuda, "pycuda.driver": driver}):
            spec.loader.exec_module(self.module)

    def test_interleaving_and_non_lifo_close_preserve_other_current_context(self):
        a, b = self.module.CudaContext(), self.module.CudaContext()
        self.assertEqual(self.stack, [])
        with b.activate():
            with a.activate():
                self.assertIs(self.stack[-1], a._ctx)
            self.assertIs(self.stack[-1], b._ctx)
            a.close()  # closing A must not pop B
            self.assertIs(self.stack[-1], b._ctx)
        self.assertEqual(self.stack, [])
        b.close()
        b.close()

    def test_exception_unwinds_and_active_close_is_rejected(self):
        a = self.module.CudaContext()
        with self.assertRaisesRegex(ValueError, "operation"):
            with a.activate():
                with self.assertRaises(RuntimeError):
                    a.close()
                raise ValueError("operation")
        self.assertEqual(self.stack, [])
        a.close()
        with self.assertRaises(RuntimeError):
            with a.activate(): pass

    def test_decorator_balances_nested_calls_and_supports_low_level_fields(self):
        bound = self.module.context_bound
        stack = self.stack
        class Owner:
            @bound
            def outer(self): return self.inner()
            @bound
            def inner(self): return stack[-1] if stack else None
        owner = Owner()
        self.assertIsNone(owner.outer())
        owner.ctx = self.module.CudaContext()
        self.assertIs(owner.outer(), owner.ctx._ctx)
        self.assertEqual(stack, [])
        owner.ctx.close()

    def test_context_manager_owns_lifetime(self):
        with self.module.CudaContext() as ctx:
            self.assertIs(self.stack[-1], ctx._ctx)
        self.assertEqual(self.stack, [])
        self.assertIsNone(ctx._ctx)


if __name__ == "__main__":
    unittest.main()