"""Host contract of PeriodicBlockStepper.step(n): one stream wait per call (mocked driver)."""
import ast
from contextlib import nullcontext
import importlib.util
from pathlib import Path
import types
from unittest.mock import Mock
import unittest

import numpy as np

from CelerisLab.config import LBMConfig
from CelerisLab.lbm.blocks import BlockLayout2D


def host_module():
    """block_stepper with pycuda/compiler replaced by mocks; no GPU context."""
    # Locate the source without importing it: importing block_stepper loads pycuda.driver,
    # which CPU-only import-boundary tests (test_build_contracts) must never see.
    path = Path(importlib.util.find_spec("CelerisLab.lbm.block_stepper").origin)
    tree = ast.parse(path.read_text())
    tree.body = [node for node in tree.body if not (
        isinstance(node, ast.Import) and any(a.name == "pycuda.driver" for a in node.names)
        or isinstance(node, ast.ImportFrom) and node.module == "cuda")]
    module = types.ModuleType("CelerisLab.lbm._test_host_block_stepper")
    module.__package__ = "CelerisLab.lbm"
    module.__dict__.update(cuda=Mock(), compiler=Mock())
    exec(compile(tree, str(path), "exec"), module.__dict__)
    return module


def backend():
    host, log = host_module(), []
    host.cuda.mem_get_info.return_value = (1 << 40, 1 << 40)
    host.cuda.mem_alloc.side_effect = lambda size: Mock(size=size)  # distinct buffers
    module = Mock()
    module.get_global.return_value = (Mock(), 20)
    module.get_function.side_effect = lambda name: Mock(
        side_effect=lambda *a, **k: log.append((name, a[:2])))
    host.compiler.load_module.return_value = module
    cfg = LBMConfig(nx=16, ny=16, viscosity=.1)
    run = host.PeriodicBlockStepper(cfg, BlockLayout2D.build(16, 16, 8, 8),
                                    context=Mock(activate=nullcontext, sm_arch="sm_70"))
    run._stream.synchronize.side_effect = lambda: log.append(("sync", ()))
    run.initialize(np.full((9, 16, 16), 1 / 9))
    log.clear()
    return run, log


class BlockStepSync(unittest.TestCase):
    def test_step_n_waits_once_and_matches_n_single_steps(self):
        many, log_many = backend()
        many.step(3)
        single, log_single = backend()
        for _ in range(3):
            single.step(1)
        names = [name for name, _ in log_many]
        self.assertEqual(names.count("sync"), 1)
        self.assertEqual(names[-1], "sync")
        self.assertEqual(names.count("BlockStep"), 3)
        # Same launch order and buffer roles as three step(1) calls.
        strip = lambda log: [(n, tuple(id(x) for x in a)) for n, a in log if n != "sync"]
        roles = lambda run, log: [(n, tuple({id(run._current): "c", id(run._spare): "s"}.get(x, "-")
                                            for x in ids)) for n, ids in strip(log)]
        self.assertEqual([n for n, _ in strip(log_many)], [n for n, _ in strip(log_single)])
        self.assertEqual(roles(many, log_many), roles(single, log_single))
        self.assertIsNot(many._current, many._spare)
        self.assertEqual((many.step_count, single.step_count), (3, 3))

    def test_step_zero_launches_and_waits_for_nothing(self):
        run, log = backend()
        run.step(0)
        self.assertEqual((log, run.step_count), ([], 0))

    def test_failure_keeps_entry_step_count(self):
        run, log = backend()
        run.step(1)
        run._stream.synchronize.side_effect = RuntimeError("device fault")
        with self.assertRaises(RuntimeError):
            run.step(2)
        self.assertEqual((run.step_count, run._status), (1, "failed"))

    def test_launch_failure_mid_batch_fails_closed_without_waiting(self):
        run, log = backend()
        calls = []
        step = run._step

        def failing(*args, **kwargs):
            calls.append(1)
            if len(calls) == 2:
                raise RuntimeError("launch failure")
            return step(*args, **kwargs)

        run._step = failing
        with self.assertRaises(RuntimeError):
            run.step(3)
        self.assertNotIn("sync", [name for name, _ in log])
        self.assertEqual((run.step_count, run._status), (0, "failed"))
        with self.assertRaises(RuntimeError):
            run.download()


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