"""Host contract of FixedPatchStepper(refinements=[]): one blocked level (mocked driver).

Per step: halo refresh -> static BFL (circles only) -> BlockStep with a null
kind, one stream wait per step(n) call, NaN spare after initialize. An upload,
launch, wait or download failure fails the instance and keeps its step count;
close retains resources for retry when context activation fails.
"""
from contextlib import nullcontext
from unittest.mock import MagicMock, Mock
import unittest

import numpy as np

from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.objects import SimObject
from CelerisLab.config import LBMConfig
from tests.unit.test_refinement_three_level_multipatch import host_module


def circle(center=(15.5, 11.5), radius=4.3):
    return SimObject(0, geometry=CircleGeometry(*center, radius), center=center, radius=radius)


def backend(objects=(), failing_kernel=None, initialize=True):
    host, log = host_module(), []
    host.armed = False  # set True to make failing_kernel raise

    def launch(name):
        def run(*a, **k):
            log.append((name, a))
            if name == failing_kernel and host.armed:
                raise RuntimeError(name + " launch failed")
        return run

    module = Mock()
    module.get_global.return_value = (Mock(), 20)
    module.get_function.side_effect = lambda name: Mock(side_effect=launch(name))
    host.compiler.build_module.return_value = module
    host.cuda.mem_alloc.side_effect = lambda size: Mock(size=size)
    host.cuda.memcpy_htod.side_effect = lambda pointer, array: log.append(("upload", (pointer, array)))
    run = host.FixedPatchStepper(LBMConfig(nx=32, ny=24, viscosity=.1),
                                 context=Mock(activate=nullcontext, sm_arch="sm_70"),
                                 refinements=[], level_blocks=[(12, 16)], level_objects=[objects])
    run._stream.synchronize.side_effect = lambda: log.append(("sync", ()))
    log.clear()
    if initialize:
        run.initialize(np.full(run._levels[0].shape, 1 / 9, np.float32))
    return run, log, host


def state():
    return np.full((9, 24, 32), 1 / 9, np.float32)


class SingleLevel(unittest.TestCase):
    def test_spare_starts_nan_and_current_holds_state(self):
        run, log, _ = backend()
        level = run._levels[0]
        uploads = {id(a[0]): a[1] for name, a in log if name == "upload"}
        self.assertTrue(np.isnan(uploads[id(level.spare)]).all())
        self.assertTrue(np.isfinite(level.layout.to_dense(uploads[id(level.current)])).all())

    def test_launch_order_per_step(self):
        for objects, order in (((), ["BlockHaloCopy", "BlockStep"]),
                               ((circle(),), ["BlockHaloCopy", "BlockCurvedBoundary", "BlockStep"])):
            with self.subTest(bodies=len(objects)):
                run, log, _ = backend(objects)
                log.clear()
                run.step(2)
                self.assertEqual([name for name, _ in log], order * 2 + ["sync"])
                step = [a for name, a in log if name == "BlockStep"][0]
                self.assertEqual(int(step[4]), 0)  # null kind: every core cell ACTIVE

    def test_circle_solids_are_not_state(self):
        run, _, _ = backend((circle(),))
        solid = run._boundaries[0].tables["solid"]
        self.assertTrue(solid.any())
        np.testing.assert_array_equal(run._state[0], ~solid)

    def test_initialize_upload_failure_fails_instance(self):
        run, _, host = backend(initialize=False)
        host.cuda.memcpy_htod.side_effect = RuntimeError("upload failed")
        with self.assertRaisesRegex(RuntimeError, "upload failed"):
            run.initialize(state())
        with self.assertRaisesRegex(RuntimeError, "failed"):
            run.step(1)

    def test_launch_failure_fails_instance_and_keeps_step_count(self):
        for kernel, objects in (("BlockHaloCopy", ()), ("BlockCurvedBoundary", (circle(),)),
                                ("BlockStep", ())):
            with self.subTest(kernel=kernel):
                run, log, host = backend(objects, failing_kernel=kernel)
                run.step(1)
                host.armed = True
                log.clear()
                with self.assertRaisesRegex(RuntimeError, kernel + " launch failed"):
                    run.step(2)
                self.assertEqual(run.step_count, 1)
                self.assertNotIn("sync", [name for name, _ in log])
                with self.assertRaisesRegex(RuntimeError, "failed"):
                    run.download()

    def test_free_memory_preflight_before_stream_and_buffers(self):
        for cfg, options in ((LBMConfig(nx=32, ny=24, viscosity=.1),
                              dict(refinements=[], level_blocks=[(12, 16)])),
                             (LBMConfig(nx=48, ny=48, viscosity=.02),  # every level counts
                              dict(refinements=[[(6, 6, 26, 26)]], level_blocks=[(4, 4)] * 2))):
            with self.subTest(depth=len(options["refinements"]) + 1):
                host = host_module()
                host.cuda.mem_get_info.return_value = (1024, 1 << 30)
                with self.assertRaisesRegex(MemoryError, "exceed free device memory"):
                    host.FixedPatchStepper(cfg, context=Mock(activate=nullcontext, sm_arch="sm_70"),
                                           **options)
                host.cuda.Stream.assert_not_called()
                host.cuda.mem_alloc.assert_not_called()
                host.compiler.build_module.assert_not_called()

    def test_wait_failure_after_a_good_step_keeps_its_count(self):
        run, _, _ = 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_final_wait_failure_fails_instance_and_keeps_step_count(self):
        run, _, _ = backend()
        run._stream.synchronize.side_effect = RuntimeError("device error")
        with self.assertRaisesRegex(RuntimeError, "device error"):
            run.step(3)
        self.assertEqual(run.step_count, 0)
        with self.assertRaisesRegex(RuntimeError, "failed"):
            run.step(1)

    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)

        def roles(run, log):  # same launch order and buffer roles as three step(1) calls
            level = run._levels[0]
            tags = {id(level.current): "c", id(level.spare): "s"}
            return [(n, tuple(tags.get(id(x), "-") for x in a)) for n, a in log if n != "sync"]
        self.assertEqual(roles(many, log_many), roles(single, log_single))
        self.assertIsNot(many._levels[0].current, many._levels[0].spare)
        self.assertEqual((many.step_count, single.step_count), (3, 3))

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

    def test_launch_failure_mid_batch_fails_without_waiting(self):
        run, log, host = backend()
        step = run._levels[0].kernel
        calls = []

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

        run._levels[0].kernel = failing
        log.clear()
        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()

    def test_initialize_rejects_bad_state_before_upload(self):
        run, log, _ = backend(initialize=False)
        good = state()
        for bad in (good[:, :-1], np.where(np.arange(32) == 3, np.nan, good), -good):
            with self.assertRaises(ValueError):
                run.initialize(bad)
        self.assertNotIn("upload", [name for name, _ in log])
        run.initialize(good)
        self.assertEqual([name for name, _ in log].count("upload"), 2)

    def test_download_failure_fails_instance(self):
        run, _, host = backend()
        host.cuda.memcpy_dtoh.side_effect = RuntimeError("read failed")
        with self.assertRaisesRegex(RuntimeError, "read failed"):
            run.download()
        for operation in (run.download, run.step):
            with self.assertRaisesRegex(RuntimeError, "failed"):
                operation()

    def test_activation_failure_retains_resources_for_retry(self):
        run, _, _ = backend()
        allocations, stream = run._allocations[:], run._stream
        unavailable = MagicMock()
        unavailable.__enter__.side_effect = RuntimeError("activation failed")
        run.ctx.activate = Mock(side_effect=[unavailable, nullcontext()])
        with self.assertRaisesRegex(RuntimeError, "activation failed"):
            run.close()
        self.assertEqual(run._status, "failed")
        self.assertEqual(run._allocations, allocations)
        self.assertIs(run._stream, stream)
        run.close()
        self.assertEqual(run._status, "closed")
        self.assertEqual(run._levels, [])
        for allocation in allocations:
            allocation.free.assert_called_once()


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