#!/usr/bin/env -S uv run --script
"""Z-Machine state inspection and debugging tool.

Loads a story file and optionally applies a Quetzal save, then displays
machine state for debugging z-machine interpreter issues.
"""

import argparse
import sys
from pathlib import Path

# Add src to path so we can import mudlib
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root / "src"))

from mudlib.zmachine.quetzal import (  # noqa: E402
    QuetzalError,
    QuetzalMismatchedFile,
    QuetzalParser,
)
from mudlib.zmachine.trivialzui import create_zui  # noqa: E402
from mudlib.zmachine.zmachine import ZMachine  # noqa: E402

# Opcode name tables for disassembly
OP2_NAMES = {
    1: "je",
    2: "jl",
    3: "jg",
    4: "dec_chk",
    5: "inc_chk",
    6: "jin",
    7: "test",
    8: "or",
    9: "and",
    10: "test_attr",
    11: "set_attr",
    12: "clear_attr",
    13: "store",
    14: "insert_obj",
    15: "loadw",
    16: "loadb",
    17: "get_prop",
    18: "get_prop_addr",
    19: "get_next_prop",
    20: "add",
    21: "sub",
    22: "mul",
    23: "div",
    24: "mod",
}

OP1_NAMES = {
    0: "jz",
    1: "get_sibling",
    2: "get_child",
    3: "get_parent",
    4: "get_prop_len",
    5: "inc",
    6: "dec",
    7: "print_addr",
    8: "call_1s",
    9: "remove_obj",
    10: "print_obj",
    11: "ret",
    12: "jump",
    13: "print_paddr",
    14: "load",
    15: "not/call_1n",
}

OP0_NAMES = {
    0: "rtrue",
    1: "rfalse",
    2: "print",
    3: "print_ret",
    4: "nop",
    5: "save",
    6: "restore",
    7: "restart",
    8: "ret_popped",
    9: "pop/catch",
    10: "quit",
    11: "new_line",
    12: "show_status",
    13: "verify",
    15: "piracy",
}

VAR_NAMES = {
    0: "call",
    1: "storew",
    2: "storeb",
    3: "put_prop",
    4: "sread",
    5: "print_char",
    6: "print_num",
    7: "random",
    8: "push",
    9: "pull",
    10: "split_window",
    11: "set_window",
    19: "output_stream",
    20: "input_stream",
    21: "sound_effect",
}


def decode_opcode_class(opcode_byte):
    """Determine opcode class from the opcode byte."""
    if opcode_byte < 0x80:
        return "2OP"
    elif opcode_byte < 0xC0:
        op_type = (opcode_byte >> 4) & 3
        if op_type == 3:
            return "0OP"
        else:
            return "1OP"
    elif opcode_byte < 0xE0:
        return "2OP"
    else:
        return "VAR"


def get_opcode_name(opcode_class, opcode_num):
    """Get the name of an opcode."""
    if opcode_class == "2OP":
        return OP2_NAMES.get(opcode_num, f"unknown_{opcode_num}")
    elif opcode_class == "1OP":
        return OP1_NAMES.get(opcode_num, f"unknown_{opcode_num}")
    elif opcode_class == "0OP":
        return OP0_NAMES.get(opcode_num, f"unknown_{opcode_num}")
    elif opcode_class == "VAR":
        return VAR_NAMES.get(opcode_num, f"unknown_{opcode_num}")
    else:
        return "unknown"


def parse_operand_types(mem, pc, opcode_byte):
    """Parse operand types without evaluating them. Returns (types, bytes_consumed)."""
    types = []
    pos = pc

    if opcode_byte < 0x80:
        # Long form 2OP
        types.append("var" if (opcode_byte & 0x40) else "small")
        types.append("var" if (opcode_byte & 0x20) else "small")
        return types, pos - pc
    elif opcode_byte < 0xC0:
        # Short form
        op_type = (opcode_byte >> 4) & 3
        if op_type == 0:
            types.append("large")
        elif op_type == 1:
            types.append("small")
        elif op_type == 2:
            types.append("var")
        # op_type == 3 means 0OP, no operands
        return types, pos - pc
    else:
        # Variable form - read types byte
        types_byte = mem[pos]
        pos += 1
        for i in range(4):
            t = (types_byte >> (6 - i * 2)) & 3
            if t == 3:
                break
            elif t == 0:
                types.append("large")
            elif t == 1:
                types.append("small")
            elif t == 2:
                types.append("var")
        return types, pos - pc


def disassemble_at(zmachine, addr, count):
    """Disassemble count instructions starting at addr."""
    mem = zmachine._mem
    pc = addr

    print(f"\n--- disassembly at 0x{addr:06x} ({count} instructions) ---")

    for _ in range(count):
        if pc >= len(mem._memory):
            print(f"  {pc:06x}  (out of bounds)")
            break

        start_pc = pc
        opcode_byte = mem[pc]
        pc += 1

        # Decode opcode class and number
        if opcode_byte < 0x80:
            # Long form 2OP
            op_class = "2OP"
            op_num = opcode_byte & 0x1F
        elif opcode_byte < 0xC0:
            # Short form
            op_type = (opcode_byte >> 4) & 3
            if op_type == 3:
                op_class = "0OP"
                op_num = opcode_byte & 0x0F
            else:
                op_class = "1OP"
                op_num = opcode_byte & 0x0F
        elif opcode_byte < 0xE0:
            # Variable form 2OP
            op_class = "2OP"
            op_num = opcode_byte & 0x1F
        else:
            # Variable form VAR
            op_class = "VAR"
            op_num = opcode_byte & 0x1F

        op_name = get_opcode_name(op_class, op_num)

        # Parse operand types
        operand_types, type_bytes = parse_operand_types(mem, pc, opcode_byte)
        pc += type_bytes

        # Skip past operand values (without reading them)
        operand_str_parts = []
        for ot in operand_types:
            if ot == "large":
                if pc + 1 < len(mem._memory):
                    val = (mem[pc] << 8) | mem[pc + 1]
                    operand_str_parts.append(f"#{val}")
                    pc += 2
            elif ot == "small":
                if pc < len(mem._memory):
                    operand_str_parts.append(f"#{mem[pc]}")
                    pc += 1
            elif ot == "var" and pc < len(mem._memory):
                operand_str_parts.append(f"V{mem[pc]:02x}")
                pc += 1

        operands_str = ", ".join(operand_str_parts)
        print(f"  {start_pc:06x}  {op_class}:{op_num:02d}  {op_name}({operands_str})")


def get_location_info(zmachine):
    """Get current location object and its contents."""
    try:
        # Global variable 0 (opcode variable 0x10) is the player location
        location_obj = zmachine._mem.read_global(0x10)

        if location_obj == 0:
            return None, []

        obj_parser = zmachine._objectparser

        # Try to get the location name
        try:
            location_name = obj_parser.get_shortname(location_obj)
        except Exception:
            # Invalid object number - return None
            return None, []

        # Get objects in this location (children)
        objects = []
        child = obj_parser.get_child(location_obj)
        while child != 0:
            try:
                child_name = obj_parser.get_shortname(child)
                objects.append(child_name)
            except Exception:
                objects.append(f"object #{child}")
            child = obj_parser.get_sibling(child)

        return (location_obj, location_name), objects

    except Exception:
        return None, []


def validate_save(zmachine, save_path):
    """Validate a save file and print diagnostic info."""
    print("\n--- save file validation ---")

    # Read the save file
    try:
        with open(save_path, "rb") as f:
            save_data = f.read()
    except OSError as e:
        print(f"ERROR: Cannot read save file: {e}")
        return False

    # Parse it
    parser = QuetzalParser(zmachine)
    try:
        parser.load_from_bytes(save_data)
    except QuetzalMismatchedFile:
        print("ERROR: Save file does not match story file")
        print("  (release number, serial, or checksum mismatch)")
        return False
    except QuetzalError as e:
        print(f"ERROR: Invalid Quetzal file: {e}")
        return False
    except Exception as e:
        print(f"ERROR: Failed to parse save file: {e}")
        return False

    # Get metadata
    metadata = parser.get_last_loaded()

    print("Save file loaded successfully")
    print(f"  release:  {metadata.get('release number', 'unknown')}")
    print(f"  serial:   {metadata.get('serial number', 'unknown')}")
    print(f"  checksum: 0x{metadata.get('checksum', 0):04x}")
    print(f"  PC:       0x{metadata.get('program counter', 0):06x}")

    # Check PC is in bounds
    pc = metadata.get("program counter", 0)
    mem_size = len(zmachine._mem._memory)
    if pc >= mem_size:
        print(f"  WARNING: PC 0x{pc:06x} is out of bounds (size: {mem_size})")
        return False

    print("  PC is within story file bounds")
    return True


def main():
    parser = argparse.ArgumentParser(
        description="Inspect z-machine state for debugging"
    )
    parser.add_argument("story", help="Path to story file (.z3/.z5/.z8)")
    parser.add_argument("--save", help="Path to Quetzal save file (.qzl)")
    parser.add_argument(
        "--disasm", type=lambda x: int(x, 0), help="Address to disassemble (hex)"
    )
    parser.add_argument(
        "--count", type=int, default=5, help="Number of instructions to disassemble"
    )
    parser.add_argument(
        "--globals", action="store_true", help="Show first 16 global variables"
    )
    args = parser.parse_args()

    story_path = Path(args.story)
    if not story_path.exists():
        print(f"ERROR: Story file not found: {story_path}")
        sys.exit(1)

    # Load story
    with open(story_path, "rb") as f:
        story_bytes = f.read()

    # Create ZMachine
    ui = create_zui()
    zmachine = ZMachine(story_bytes, ui)

    # Load save if provided
    if args.save:
        save_path = Path(args.save)
        if not save_path.exists():
            print(f"ERROR: Save file not found: {save_path}")
            sys.exit(1)

        if not validate_save(zmachine, save_path):
            sys.exit(1)

    # Display state
    print("\nzmachine state")
    print("==============")
    print(f"version:    {zmachine._mem.version}")
    print(f"story size: {len(story_bytes)} bytes")
    print(f"PC:         0x{zmachine._opdecoder.program_counter:06x}")
    print(f"stack:      {len(zmachine._stackmanager._call_stack) - 1} frames")

    # Location info
    location, objects = get_location_info(zmachine)
    if location:
        loc_obj, loc_name = location
        print(f"\nlocation:   {loc_name} (#{loc_obj})")
        if objects:
            print("objects:    " + ", ".join(objects))
        else:
            print("objects:    (none)")
    else:
        print("\nlocation:   (unknown)")

    # Global variables
    if args.globals:
        print("\nglobals (first 16):")
        for i in range(16):
            global_val = zmachine._mem.read_global(0x10 + i)
            print(f"  G{i:02d} (V{0x10 + i:02x}): 0x{global_val:04x} ({global_val})")

    # Disassembly
    if args.disasm is not None:
        disassemble_at(zmachine, args.disasm, args.count)
    else:
        # Disassemble instructions at PC
        disassemble_at(zmachine, zmachine._opdecoder.program_counter, args.count)


if __name__ == "__main__":
    main()