"""Build a 48K ZX Spectrum .SNA snapshot (and raw .SCR) from a screen image.

A .SNA bakes the whole 48K RAM plus the CPU state.  We put the 6912-byte screen
at $4000, a 3-byte idle stub (DI; JR $) at $8000, and set SP so the loader's
RETN jumps to the stub -- the ULA then displays the screen forever while the CPU
idles, so the picture appears the instant MAME loads the file.

.SNA 27-byte header: I, HL', DE', BC', AF', HL, DE, BC, IY, IX, IFF2, R, AF, SP,
IM, border; followed by 49152 bytes of RAM ($4000-$FFFF).
"""

from __future__ import annotations

import struct

RAM_BASE = 0x4000
RAM_SIZE = 0xC000          # 48K ($4000-$FFFF)
SCREEN_LEN = 6912          # 6144 bitmap + 768 attributes
STUB_ADDR = 0x8000         # idle loop  DI; JR $
STUB = bytes([0xF3, 0x18, 0xFE])
SP_ADDR = 0xFF00           # stack holds the stub address for the loader's RETN


def build_sna(screen: bytes, border: int = 0) -> bytes:
    if len(screen) != SCREEN_LEN:
        raise ValueError(f"screen must be {SCREEN_LEN} bytes, got {len(screen)}")
    ram = bytearray(RAM_SIZE)
    ram[0x0000:SCREEN_LEN] = screen                    # $4000 screen
    off = STUB_ADDR - RAM_BASE
    ram[off:off + len(STUB)] = STUB                    # $8000 idle stub
    sp = SP_ADDR - RAM_BASE                            # return address for RETN
    ram[sp] = STUB_ADDR & 0xFF
    ram[sp + 1] = (STUB_ADDR >> 8) & 0xFF

    header = bytearray(27)
    header[0x00] = 0x3F                                # I
    # HL',DE',BC',AF',HL,DE,BC,IY,IX all zero
    header[0x13] = 0x00                                # IFF2 = 0 (interrupts off)
    header[0x14] = 0x00                                # R
    struct.pack_into("<H", header, 0x17, SP_ADDR)      # SP
    header[0x19] = 0x01                                # IM 1
    header[0x1A] = border & 0x07                       # border colour
    return bytes(header) + bytes(ram)


def build_scr(screen: bytes) -> bytes:
    """The standard 6912-byte ZX Spectrum screen file (raw $4000 dump)."""
    if len(screen) != SCREEN_LEN:
        raise ValueError(f"screen must be {SCREEN_LEN} bytes, got {len(screen)}")
    return bytes(screen)