"""CPU mocks execute production remap lifecycle; no GPU numerical claim."""
from dataclasses import replace
import unittest
from unittest.mock import patch

import numpy as np

from tests.unit.test_refinement_three_level_multipatch import mock_backend
from tests.unit.test_synchronized_remap import A, B, fields, topology, remap


def ready_backend():
    run, module, allocations = mock_backend(children=A)
    storage = {}
    module.cuda.memcpy_htod.side_effect = lambda ptr, a: storage.__setitem__(id(ptr), np.array(a, copy=True))
    module.cuda.memcpy_dtoh.side_effect = lambda a, ptr: np.copyto(a, storage[id(ptr)])
    run.initialize(*fields(topology(A), np.float32))
    run._step_count = 7
    return run, module, allocations, storage


class RuntimeChildRemap(unittest.TestCase):
    def test_current_snapshot_identity_and_count_continuation(self):
        run, module, allocations, storage = ready_backend()
        # Make current state distinct from initialization and spare, as after a step.
        for level in run._levels:
            storage[id(level.current)] *= np.float32(1.01)
        before = run.download()
        saved = {key: value.copy() for key, value in storage.items()}
        old_allocations = list(allocations)
        module.cuda.memcpy_dtoh.reset_mock()
        clone = run.remap_child(A)
        self.assertEqual(module.cuda.memcpy_dtoh.call_count, 3)
        for a, b in zip(clone.download(), before):
            np.testing.assert_array_equal(a, b)
        self.assertEqual((run.step_count, clone.step_count), (7, 7))
        for key, value in saved.items():
            np.testing.assert_array_equal(storage[key], value)
        for allocation in old_allocations:
            allocation.free.assert_not_called()
        clone.step()
        self.assertEqual((run.step_count, clone.step_count), (7, 8))
        run.step()
        self.assertEqual(run.step_count, 8)
        clone.close(); run.close()

    def test_shift_uses_cpu_remap_and_resets_both_buffers(self):
        run, module, allocations, storage = ready_backend()
        before = run.download()
        expected = remap(topology(A), topology(B), before).g
        clone = run.remap_child(B)
        for level, want, got in zip(clone._levels, expected, clone.download()):
            np.testing.assert_array_equal(got, want)
            for ptr in (level.current, level.spare):
                np.testing.assert_array_equal(storage[id(ptr)], level.layout.to_blocked(want))
        for a, b in zip(run.download(), before):
            np.testing.assert_array_equal(a, b)
        clone.initialize(*expected)
        self.assertEqual(clone.step_count, 0)  # Existing public initialization semantics.
        clone.close(); run.close()

    def test_invalid_state_geometry_and_body_reject_before_reads_allocations(self):
        run, module, allocations, storage = ready_backend()
        start = len(allocations)
        module.cuda.memcpy_dtoh.reset_mock()
        for status in ('new', 'failed', 'closed', 'advancing'):
            run._status = status
            with self.assertRaises(RuntimeError): run.remap_child(B)
        run._status = 'ready'
        for bad in ([], [(40,16,48,40)], [(16,16,32,24),(16,27,32,55)]):
            with self.assertRaises(ValueError): run.remap_child(bad)
        hosts, pairs = run._boundary_hosts, run.hierarchy.pairs
        for attribute, value in (('_wall_aware', True),
                                 ('_boundary_hosts', [*hosts[:2], object()]),        # child boundary
                                 ('_boundary_hosts', [hosts[0], object(), hosts[2]]),  # fine boundary
                                 ('hierarchy', replace(run.hierarchy, pairs=(pairs[0], None)))):
            with patch.object(run, attribute, value):
                with self.assertRaises(ValueError): run.remap_child(B)
        self.assertEqual(len(allocations), start)
        module.cuda.memcpy_dtoh.assert_not_called()
        run.close()

    def test_candidate_construction_or_upload_failure_preserves_source(self):
        for failure in ('construction', 'initialize'):
            run, module, allocations, storage = ready_backend()
            before = run.download()
            old_allocations = list(allocations)
            if failure == 'construction':
                operation = patch.object(module.compiler, 'build_module', side_effect=RuntimeError('candidate compile'))
            else:
                initialize = type(run).initialize
                def fail(candidate, *arrays):
                    upload = module.cuda.memcpy_htod.side_effect
                    calls = 0
                    def partial_upload(pointer, array):
                        nonlocal calls
                        calls += 1
                        if calls == 2:
                            raise RuntimeError('candidate upload')
                        upload(pointer, array)
                    with patch.object(module.cuda, 'memcpy_htod', side_effect=partial_upload):
                        initialize(candidate, *arrays)
                operation = patch.object(type(run), 'initialize', fail)
            with operation, self.assertRaisesRegex(RuntimeError, 'candidate'):
                run.remap_child(B)
            self.assertEqual((run._status, run.step_count), ('ready', 7))
            for a, b in zip(run.download(), before): np.testing.assert_array_equal(a, b)
            for allocation in old_allocations: allocation.free.assert_not_called()
            for allocation in allocations[len(old_allocations):]: allocation.free.assert_called_once()
            run.close()

    def test_old_sync_or_read_failure_marks_only_old_failed(self):
        for operation in ('sync', 'read'):
            run, module, allocations, storage = ready_backend()
            start = len(allocations)
            target = run._stream.synchronize if operation == 'sync' else module.cuda.memcpy_dtoh
            target.side_effect = RuntimeError('old device failure')
            with self.assertRaisesRegex(RuntimeError, 'old device'):
                run.remap_child(B)
            self.assertEqual(run._status, 'failed')
            self.assertEqual(len(allocations), start)
            target.side_effect = None
            run.close()