"""Build an Amstrad CPC .SNA snapshot (CPCEMU v1) from a screen image.

MAME loads a .sna by restoring the Z80, Gate Array (mode + 16-pen palette),
CRTC and 64K RAM, so the picture appears instantly with no loader.  We bake the
16K screen at &C000, set the Gate Array palette + mode, program a standard 80x200
CRTC screen, and point the CPU at a tiny ``DI: JR $`` idle stub -- the CRTC then
DMAs the screen forever while the Z80 idles.

Header layout (offsets) per MAME's amstrad_handle_snapshot:
  0x00 "MV - SNA" signature   0x10 version
  0x11 AF 0x13 BC 0x15 DE 0x17 HL 0x19 R 0x1a I 0x1b IFF1 0x1c IFF2
  0x1d IX 0x1f IY 0x21 SP 0x23 PC 0x25 IM 0x26 AF' .. 0x2c HL'
  0x2e GA selected pen   0x2f..0x3f 17 ink numbers (16 pens + border)
  0x40 GA multi-config (mode/rom)   0x41 RAM config
  0x42 CRTC selected reg  0x43..0x54 18 CRTC regs   0x55 upper ROM
  0x56..0x58 PPI A/B/C  0x59 PPI control  0x5a PSG reg  0x5b..0x6a PSG regs
  0x6b memory size (KB)   0x100 RAM dump
"""
from __future__ import annotations

import struct

SCREEN_BASE = 0xC000
SCREEN_LEN = 0x4000           # 16K
STUB_ADDR = 0x8000            # idle loop  DI; JR $  (central RAM, always present)
STUB = bytes([0xF3, 0x18, 0xFE])
RAM_SIZE = 0x10000            # 64K

# Standard CPC 50Hz CRTC register set for a 40x25-char (80x200 byte) screen at
# &C000 (R12/R13 = 0x30/0x00).
CRTC = [0x3F, 0x28, 0x2E, 0x8E, 0x26, 0x00, 0x19, 0x1E,
        0x00, 0x07, 0x00, 0x00, 0x30, 0x00, 0x00, 0x00, 0x00, 0x00]


def screen_offset(y: int, bx: int) -> int:
    """Offset within the 16K screen for scan line ``y`` (0-199), byte column
    ``bx`` (0-79) -- the CPC's CRTC interleave (8 lines per char row)."""
    return (y % 8) * 0x800 + (y // 8) * 80 + bx


def build_sna(screen: bytes, inks, mode: int, border: int = 20) -> bytes:
    """screen: 16384 bytes (&C000 RAM); inks: 16 hardware ink numbers (pens);
    mode: 0/1/2; border: hardware ink number for the border pen."""
    if len(screen) != SCREEN_LEN:
        raise ValueError(f"screen must be {SCREEN_LEN} bytes, got {len(screen)}")
    if not 1 <= len(inks) <= 16:
        raise ValueError(f"need 1-16 ink numbers, got {len(inks)}")
    # the Gate Array always has 16 pen slots; modes 1/2 use only the first 4/2.
    inks = list(inks) + [border] * (16 - len(inks))

    ram = bytearray(RAM_SIZE)
    ram[STUB_ADDR:STUB_ADDR + len(STUB)] = STUB
    ram[SCREEN_BASE:SCREEN_BASE + SCREEN_LEN] = screen

    h = bytearray(0x100)
    h[0x00:0x10] = b"MV - SNAPSHOT V1"          # signature (MAME checks "MV - SNA")
    h[0x10] = 1                                   # version
    # Z80: everything 0 except SP/PC/IM; DI (IFF=0) and HALT-free idle stub.
    struct.pack_into("<H", h, 0x21, STUB_ADDR)    # SP
    struct.pack_into("<H", h, 0x23, STUB_ADDR)    # PC
    h[0x25] = 1                                    # IM 1
    # Gate Array palette: 16 pens + border, each a 5-bit hardware ink number.
    h[0x2e] = 0                                    # selected pen
    for pen in range(16):
        h[0x2f + pen] = inks[pen] & 0x1F
    h[0x2f + 16] = border & 0x1F                   # border ink
    h[0x40] = (mode & 0x03) | 0x0C                 # mode + both ROMs disabled (RAM)
    h[0x41] = 0x00                                 # RAM config 0 (standard 64K)
    h[0x42] = 0                                    # CRTC selected reg
    for i, v in enumerate(CRTC):
        h[0x43 + i] = v & 0xFF
    h[0x55] = 0                                    # upper ROM
    h[0x59] = 0x82                                 # PPI control (conventional)
    h[0x5b + 7] = 0x3F                             # PSG mixer: all channels off (silent)
    struct.pack_into("<H", h, 0x6b, 64)           # 64K dump

    return bytes(h) + bytes(ram)