"""
Tests for QuetzalWriter._generate_stks_chunk() serialization.

The Stks chunk serializes the Z-machine call stack. Each frame has:
- Bytes 0-2: return_pc (24-bit big-endian) — caller's resume PC
- Byte 3: flags (bits 0-3 = num local vars, bit 4 = discard result)
- Byte 4: varnum (which variable gets return value)
- Byte 5: argflag (bitmask of supplied arguments)
- Bytes 6-7: eval_stack_size (16-bit big-endian)
- Next (num_local_vars * 2) bytes: local variables
- Next (eval_stack_size * 2) bytes: evaluation stack values

All multi-byte values are big-endian. Bottom frame has return_pc=0.

Field mapping to runtime:
- return_pc for frame N → stored as program_counter on frame N-1 (the caller)
- varnum → stored as return_addr on the frame (the store variable)
"""

from unittest import TestCase

from mudlib.zmachine.quetzal import QuetzalWriter
from mudlib.zmachine.zstackmanager import ZRoutine, ZStackManager


class MockMemory:
    """Mock memory for testing."""

    def __init__(self):
        self.version = 5


class MockZMachine:
    """Mock z-machine with stack manager."""

    def __init__(self):
        self._mem = MockMemory()
        self._stackmanager = ZStackManager(self._mem)


class QuetzalStksTests(TestCase):
    """Test suite for Stks chunk generation."""

    def setUp(self):
        self.zmachine = MockZMachine()
        self.writer = QuetzalWriter(self.zmachine)

    def test_empty_call_stack_generates_empty_chunk(self):
        """With only the sentinel bottom, should generate empty bytes."""
        # Call stack has only ZStackBottom sentinel
        chunk = self.writer._generate_stks_chunk()
        self.assertEqual(chunk, b"")

    def test_single_frame_serialization(self):
        """Single routine frame should serialize correctly."""
        # Set up caller resume PC on ZStackBottom
        sentinel = self.zmachine._stackmanager._call_stack[0]
        sentinel.program_counter = 0x4200

        # Create a routine with return_addr = store variable (varnum 5)
        routine = ZRoutine(
            start_addr=0x5000,
            return_addr=5,
            zmem=self._mem,
            args=[],
            local_vars=[0x1234, 0x5678, 0xABCD],
            stack=[0x1111, 0x2222],
        )
        self.zmachine._stackmanager.push_routine(routine)

        chunk = self.writer._generate_stks_chunk()

        # return_pc comes from ZStackBottom.program_counter (0x4200)
        # varnum is frame.return_addr (5)
        expected = bytes(
            [
                0x00,
                0x42,
                0x00,  # return_pc (from caller's program_counter)
                0x03,  # flags (3 local vars)
                0x05,  # varnum (store variable)
                0x00,  # argflag
                0x00,
                0x02,  # eval_stack_size = 2
                0x12,
                0x34,  # local_vars[0]
                0x56,
                0x78,  # local_vars[1]
                0xAB,
                0xCD,  # local_vars[2]
                0x11,
                0x11,  # stack[0]
                0x22,
                0x22,  # stack[1]
            ]
        )
        self.assertEqual(chunk, expected)

    def test_multiple_frames_serialization(self):
        """Multiple nested frames should serialize in order."""
        sentinel = self.zmachine._stackmanager._call_stack[0]
        sentinel.program_counter = 0  # main routine has no caller

        # Frame 1: outer routine (varnum=0 means push result to stack)
        routine1 = ZRoutine(
            start_addr=0x1000,
            return_addr=0,
            zmem=self._mem,
            args=[],
            local_vars=[0x0001],
            stack=[],
        )
        self.zmachine._stackmanager.push_routine(routine1)
        # Set routine1's resume PC (where to go after frame2 returns)
        routine1.program_counter = 0x1050

        # Frame 2: inner routine (varnum=3)
        routine2 = ZRoutine(
            start_addr=0x2000,
            return_addr=3,
            zmem=self._mem,
            args=[],
            local_vars=[0x0002, 0x0003],
            stack=[0xAAAA],
        )
        self.zmachine._stackmanager.push_routine(routine2)

        chunk = self.writer._generate_stks_chunk()

        # Frame 1: return_pc from sentinel.pc (0), varnum=0
        frame1 = bytes(
            [
                0x00,
                0x00,
                0x00,  # return_pc = 0 (from sentinel)
                0x01,  # flags (1 local var)
                0x00,  # varnum = 0 (push to stack)
                0x00,  # argflag
                0x00,
                0x00,  # eval_stack_size = 0
                0x00,
                0x01,  # local_vars[0]
            ]
        )
        # Frame 2: return_pc from routine1.pc (0x1050), varnum=3
        frame2 = bytes(
            [
                0x00,
                0x10,
                0x50,  # return_pc (from routine1.program_counter)
                0x02,  # flags (2 local vars)
                0x03,  # varnum = 3
                0x00,  # argflag
                0x00,
                0x01,  # eval_stack_size = 1
                0x00,
                0x02,  # local_vars[0]
                0x00,
                0x03,  # local_vars[1]
                0xAA,
                0xAA,  # stack[0]
            ]
        )
        expected = frame1 + frame2
        self.assertEqual(chunk, expected)

    def test_frame_with_no_locals_or_stack(self):
        """Frame with no local vars or stack values should serialize correctly."""
        sentinel = self.zmachine._stackmanager._call_stack[0]
        sentinel.program_counter = 0x2500

        routine = ZRoutine(
            start_addr=0x3000,
            return_addr=1,
            zmem=self._mem,
            args=[],
            local_vars=[],
            stack=[],
        )
        self.zmachine._stackmanager.push_routine(routine)

        chunk = self.writer._generate_stks_chunk()

        expected = bytes(
            [
                0x00,
                0x25,
                0x00,  # return_pc (from sentinel.program_counter)
                0x00,  # flags (0 local vars)
                0x01,  # varnum = 1
                0x00,  # argflag
                0x00,
                0x00,  # eval_stack_size = 0
            ]
        )
        self.assertEqual(chunk, expected)

    def test_discard_result_sets_flags_bit4(self):
        """Frame with return_addr=None should set bit 4 in flags byte."""
        sentinel = self.zmachine._stackmanager._call_stack[0]
        sentinel.program_counter = 0

        routine = ZRoutine(
            start_addr=0x1000,
            return_addr=None,
            zmem=self._mem,
            args=[],
            local_vars=[0x0001, 0x0002],
            stack=[],
        )
        self.zmachine._stackmanager.push_routine(routine)

        chunk = self.writer._generate_stks_chunk()

        # flags = 0x02 (2 locals) | 0x10 (discard) = 0x12
        expected = bytes(
            [
                0x00,
                0x00,
                0x00,  # return_pc
                0x12,  # flags (2 locals + discard bit)
                0x00,  # varnum (0 when discarding)
                0x00,  # argflag
                0x00,
                0x00,  # eval_stack_size = 0
                0x00,
                0x01,  # local_vars[0]
                0x00,
                0x02,  # local_vars[1]
            ]
        )
        self.assertEqual(chunk, expected)

    def test_round_trip_with_parser(self):
        """Generated stks bytes should parse back identically."""
        from mudlib.zmachine.quetzal import QuetzalParser

        sentinel = self.zmachine._stackmanager._call_stack[0]
        sentinel.program_counter = 0  # main routine caller

        # Create a complex stack state
        routine1 = ZRoutine(
            start_addr=0x1000,
            return_addr=0,
            zmem=self._mem,
            args=[],
            local_vars=[0x1111, 0x2222, 0x3333],
            stack=[0xAAAA, 0xBBBB],
        )
        self.zmachine._stackmanager.push_routine(routine1)
        routine1.program_counter = 0x1234  # resume PC for after routine2

        routine2 = ZRoutine(
            start_addr=0x2000,
            return_addr=5,
            zmem=self._mem,
            args=[],
            local_vars=[0x4444],
            stack=[0xCCCC, 0xDDDD, 0xEEEE],
        )
        self.zmachine._stackmanager.push_routine(routine2)

        # Generate the stks chunk
        stks_bytes = self.writer._generate_stks_chunk()

        # Parse it back
        parser = QuetzalParser(self.zmachine)
        parser._parse_stks(stks_bytes)

        # Verify the stack was reconstructed correctly
        # Parser creates a new stack manager, skip bottom sentinel
        call_stack = self.zmachine._stackmanager._call_stack
        frames = call_stack[1:]
        self.assertEqual(len(frames), 2)

        # Check frame 1: return_addr = varnum, caller PC on sentinel
        assert isinstance(frames[0], ZRoutine)
        self.assertEqual(frames[0].return_addr, 0)
        self.assertEqual(frames[0].local_vars[:3], [0x1111, 0x2222, 0x3333])
        self.assertEqual(frames[0].stack, [0xAAAA, 0xBBBB])
        # Sentinel should have frame1's return_pc (0)
        self.assertEqual(call_stack[0].program_counter, 0)

        # Check frame 2: return_addr = varnum, caller PC on frame1
        assert isinstance(frames[1], ZRoutine)
        self.assertEqual(frames[1].return_addr, 5)
        self.assertEqual(frames[1].local_vars[:1], [0x4444])
        self.assertEqual(frames[1].stack, [0xCCCC, 0xDDDD, 0xEEEE])
        # Frame1 should have frame2's return_pc (0x1234)
        self.assertEqual(frames[0].program_counter, 0x1234)

    def test_parse_return_pc_goes_to_caller(self):
        """Parser should put return_pc on the caller frame's program_counter."""
        from mudlib.zmachine.quetzal import QuetzalParser

        # Construct a minimal stack frame with return_pc=0x123456
        stks_bytes = bytes(
            [
                0x12,
                0x34,
                0x56,  # return_pc = 0x123456
                0x00,  # flags (0 local vars)
                0x07,  # varnum = 7
                0x00,  # argflag
                0x00,
                0x00,  # eval_stack_size = 0
            ]
        )

        parser = QuetzalParser(self.zmachine)
        parser._parse_stks(stks_bytes)

        call_stack = self.zmachine._stackmanager._call_stack
        frames = call_stack[1:]  # skip sentinel
        self.assertEqual(len(frames), 1)

        # return_pc goes to caller (sentinel) program_counter
        self.assertEqual(call_stack[0].program_counter, 0x123456)
        # varnum goes to frame's return_addr
        assert isinstance(frames[0], ZRoutine)
        self.assertEqual(frames[0].return_addr, 7)

    def test_parse_discard_bit_restores_none(self):
        """Parser should set return_addr=None when flags bit 4 is set."""
        from mudlib.zmachine.quetzal import QuetzalParser

        # flags = 0x12 = 2 locals + discard bit
        stks_bytes = bytes(
            [
                0x00,
                0x00,
                0x00,  # return_pc = 0
                0x12,  # flags (2 locals + discard)
                0x00,  # varnum (ignored when discarding)
                0x00,  # argflag
                0x00,
                0x00,  # eval_stack_size = 0
                0x00,
                0x01,  # local_vars[0]
                0x00,
                0x02,  # local_vars[1]
            ]
        )

        parser = QuetzalParser(self.zmachine)
        parser._parse_stks(stks_bytes)

        frames = self.zmachine._stackmanager._call_stack[1:]
        self.assertEqual(len(frames), 1)
        self.assertIsNone(frames[0].return_addr)
        self.assertEqual(frames[0].local_vars[:2], [1, 2])

    @property
    def _mem(self):
        """Helper to get mock memory."""
        return self.zmachine._mem