"""Tests for Z-machine save_undo / restore_undo opcodes (V5+, EXT:9/10)."""

from unittest import TestCase
from unittest.mock import Mock

from mudlib.zmachine.zcpu import ZCpu
from mudlib.zmachine.zmemory import ZMemory
from mudlib.zmachine.zstackmanager import ZStackManager


class MockOpDecoder:
    """Mock opcode decoder for undo tests."""

    def __init__(self):
        self.program_counter = 0x800
        self.store_address = None
        self.branch_condition = True
        self.branch_offset = 2

    def get_store_address(self):
        return self.store_address

    def get_branch_offset(self):
        return (self.branch_condition, self.branch_offset)


class MockUI:
    def __init__(self):
        self.screen = Mock()
        self.keyboard_input = Mock()
        self.filesystem = Mock()


def make_v8_story(static_start=0x0800, globals_start=0x0400):
    """Create a minimal V8 story with proper memory layout."""
    size = max(static_start + 512, 2048)
    story = bytearray(size)
    story[0] = 8  # V8
    story[0x04] = (static_start >> 8) & 0xFF  # high memory start
    story[0x05] = static_start & 0xFF
    story[0x0C] = (globals_start >> 8) & 0xFF
    story[0x0D] = globals_start & 0xFF
    story[0x0E] = (static_start >> 8) & 0xFF
    story[0x0F] = static_start & 0xFF
    return story


class SaveUndoTests(TestCase):
    """Tests for op_save_undo and op_restore_undo."""

    def setUp(self):
        story = make_v8_story()
        self.memory = ZMemory(bytes(story))
        self.stack = ZStackManager(self.memory)
        self.decoder = MockOpDecoder()
        self.ui = MockUI()

        self.cpu = ZCpu(
            self.memory,
            self.decoder,
            self.stack,
            Mock(),  # objects
            Mock(),  # string
            Mock(),  # stream manager
            self.ui,
            Mock(),  # lexer
            zmachine=None,
        )

    def test_save_undo_stores_1(self):
        """save_undo stores 1 (success), not -1 (not available)."""
        self.decoder.store_address = 0x10
        self.cpu.op_save_undo()
        self.assertEqual(self.memory.read_global(0x10), 1)

    def test_restore_undo_no_snapshot_stores_0(self):
        """restore_undo with no prior save stores 0."""
        self.decoder.store_address = 0x10
        self.cpu.op_restore_undo()
        self.assertEqual(self.memory.read_global(0x10), 0)

    def test_save_then_restore_stores_2(self):
        """After save + restore, save_undo's store location holds 2."""
        # save_undo stores result in global var 0x10
        self.decoder.store_address = 0x10
        self.decoder.program_counter = 0x900
        self.cpu.op_save_undo()
        self.assertEqual(self.memory.read_global(0x10), 1)

        # restore — use a different store address for restore_undo itself
        self.decoder.store_address = 0x11
        self.decoder.program_counter = 0xA00
        self.cpu.op_restore_undo()

        # save_undo's store var (0x10) should hold 2
        self.assertEqual(self.memory.read_global(0x10), 2)

    def test_restore_reverts_dynamic_memory(self):
        """Memory changes after save_undo are reverted by restore_undo."""
        self.decoder.store_address = 0x10
        self.decoder.program_counter = 0x900

        # Set initial memory state in dynamic region
        self.memory[0x100] = 0xAA
        self.memory[0x200] = 0xBB

        self.cpu.op_save_undo()

        # Modify memory after save
        self.memory[0x100] = 0x11
        self.memory[0x200] = 0x22
        self.memory[0x300] = 0x33

        # Restore
        self.decoder.store_address = 0x11
        self.cpu.op_restore_undo()

        # Memory should be reverted to save-time state
        self.assertEqual(self.memory[0x100], 0xAA)
        self.assertEqual(self.memory[0x200], 0xBB)
        self.assertEqual(self.memory[0x300], 0x00)

    def test_restore_reverts_program_counter(self):
        """PC is restored to save_undo's save point."""
        self.decoder.store_address = 0x10
        self.decoder.program_counter = 0x900
        self.cpu.op_save_undo()

        # Move PC forward
        self.decoder.program_counter = 0xB00

        # Restore
        self.decoder.store_address = 0x11
        self.cpu.op_restore_undo()

        self.assertEqual(self.decoder.program_counter, 0x900)

    def test_restore_reverts_call_stack(self):
        """Call stack changes after save_undo are reverted."""
        self.decoder.store_address = 0x10
        self.decoder.program_counter = 0x900

        # Set up a routine header at 0x500: 2 local vars (V8 = zero-init)
        self.memory._memory[0x500] = 2

        # Start a routine before saving
        self.stack.start_routine(0x500, 0x01, 0x900, [42, 99])
        self.assertEqual(len(self.stack._call_stack), 2)

        self.cpu.op_save_undo()

        # Push another routine after save
        self.memory._memory[0x600] = 1
        self.stack.start_routine(0x600, 0x02, 0xA00, [7])
        self.assertEqual(len(self.stack._call_stack), 3)

        # Restore
        self.decoder.store_address = 0x11
        self.cpu.op_restore_undo()

        # Should be back to 2 frames
        self.assertEqual(len(self.stack._call_stack), 2)

    def test_multiple_saves_keeps_latest(self):
        """Second save_undo overwrites the first snapshot."""
        self.decoder.store_address = 0x10

        # First save
        self.memory[0x100] = 0xAA
        self.decoder.program_counter = 0x900
        self.cpu.op_save_undo()

        # Modify and save again
        self.memory[0x100] = 0xBB
        self.decoder.program_counter = 0x950
        self.cpu.op_save_undo()

        # Modify again
        self.memory[0x100] = 0xCC

        # Restore — should go to second save (0xBB), not first (0xAA)
        self.decoder.store_address = 0x11
        self.cpu.op_restore_undo()

        self.assertEqual(self.memory[0x100], 0xBB)
        self.assertEqual(self.decoder.program_counter, 0x950)

    def test_snapshot_is_deep_copy_memory(self):
        """Modifying memory after save doesn't corrupt the snapshot."""
        self.decoder.store_address = 0x10
        self.decoder.program_counter = 0x900

        self.memory[0x100] = 0xAA
        self.cpu.op_save_undo()

        # Overwrite the same address
        self.memory[0x100] = 0xFF

        # Restore
        self.decoder.store_address = 0x11
        self.cpu.op_restore_undo()

        # Should be AA, not FF
        self.assertEqual(self.memory[0x100], 0xAA)

    def test_snapshot_is_deep_copy_stack(self):
        """Modifying stack frames after save doesn't corrupt the snapshot."""
        self.decoder.store_address = 0x10
        self.decoder.program_counter = 0x900

        # Start a routine with known local vars
        self.memory._memory[0x500] = 2
        self.stack.start_routine(0x500, 0x01, 0x900, [42, 99])
        self.stack.push_stack(555)

        self.cpu.op_save_undo()

        # Mutate the live stack frame
        self.stack.set_local_variable(0, 0)
        self.stack.push_stack(999)

        # Restore
        self.decoder.store_address = 0x11
        self.cpu.op_restore_undo()

        # Local var 0 should be restored to 42
        self.assertEqual(self.stack.get_local_variable(0), 42)
        # Eval stack should have the original 555 (not 999)
        self.assertEqual(self.stack.pop_stack(), 555)

    def test_restore_consumes_snapshot(self):
        """After restore, the snapshot is consumed (no double undo)."""
        self.decoder.store_address = 0x10
        self.decoder.program_counter = 0x900
        self.cpu.op_save_undo()

        # First restore succeeds
        self.decoder.store_address = 0x11
        self.cpu.op_restore_undo()
        self.assertEqual(self.memory.read_global(0x10), 2)

        # Second restore fails (no snapshot)
        self.decoder.store_address = 0x12
        self.cpu.op_restore_undo()
        self.assertEqual(self.memory.read_global(0x12), 0)
