"""Build a Sega Master System .sms cartridge: a Z80 VDP-setup program + data.

The Z80 code (org $0000) programs the VDP for mode 4 (256x192), uploads the tile
patterns to VRAM $0000, the name table to $3800 and the palette to CRAM, turns
the display on, then idles.  A valid "TMR SEGA" header + checksum is written at
$7FF0 so the export BIOS accepts the cartridge.
"""
from __future__ import annotations

from .z80 import Asm

ROM_SIZE = 0x8000              # 32 KB (maps to $0000-$7FFF, no mapper needed)
PORT_CTRL = 0xBF
PORT_DATA = 0xBE

# VDP registers 0-10 for a plain 256x192 mode-4 background screen.  Reg 1 starts
# with the display OFF ($80); it is set to $C0 (display on) at the end.
VDP_REGS = [0x04, 0x80, 0xFF, 0xFF, 0xFF, 0xFF, 0xFB, 0x00, 0x00, 0x00, 0xFF]


def _code(reglen, patlen, ntlen) -> bytes:
    a = Asm(0x0000)
    a.di()
    a.im1()
    a.ld_sp(0xDFF0)

    # --- VDP register init: write VDP_REGS[i] then $80|i ---
    a.ld_hl("REGDATA")
    a.ld_b(len(VDP_REGS))
    a.ld_c(0x00)
    a.label("rinit")
    a.ld_a_hl(); a.out_a(PORT_CTRL)      # value
    a.ld_a_c(); a.or_n(0x80); a.out_a(PORT_CTRL)   # $80 | reg
    a.inc_hl(); a.inc_c()
    a.djnz("rinit")

    # --- upload patterns to VRAM $0000 ---
    a.xor_a(); a.out_a(PORT_CTRL)
    a.ld_a(0x40); a.out_a(PORT_CTRL)     # $0000 | write
    a.ld_hl("PATDATA")
    a.ld_bc(patlen)
    a.label("ptile")
    a.ld_a_hl(); a.out_a(PORT_DATA)
    a.inc_hl(); a.dec_bc()
    a.ld_a_b(); a.or_c(); a.jp_nz("ptile")

    # --- upload name table to VRAM $3800 ---
    a.xor_a(); a.out_a(PORT_CTRL)
    a.ld_a(0x78); a.out_a(PORT_CTRL)     # $3800 | write ($38 | $40)
    a.ld_hl("NTDATA")
    a.ld_bc(ntlen)
    a.label("pnt")
    a.ld_a_hl(); a.out_a(PORT_DATA)
    a.inc_hl(); a.dec_bc()
    a.ld_a_b(); a.or_c(); a.jp_nz("pnt")

    # --- upload palette to CRAM 0 ---
    a.xor_a(); a.out_a(PORT_CTRL)
    a.ld_a(0xC0); a.out_a(PORT_CTRL)     # CRAM write
    a.ld_hl("PALDATA")
    a.ld_b(0x20)                         # 32 colours
    a.label("ppal")
    a.ld_a_hl(); a.out_a(PORT_DATA)
    a.inc_hl()
    a.djnz("ppal")

    # --- disable sprites: write $D0 (list terminator) to SAT Y[0] at $3F00 ---
    a.xor_a(); a.out_a(PORT_CTRL)
    a.ld_a(0x7F); a.out_a(PORT_CTRL)     # $3F00 | write
    a.ld_a(0xD0); a.out_a(PORT_DATA)

    # --- display on (reg 1 = $C0) ---
    a.ld_a(0xC0); a.out_a(PORT_CTRL)
    a.ld_a(0x81); a.out_a(PORT_CTRL)
    a.label("hang")
    a.jp("hang")

    # data labels live right after the code, in this order
    base = len(a.code)
    a.set_label("REGDATA", base)
    a.set_label("PATDATA", base + reglen)
    a.set_label("NTDATA", base + reglen + patlen)
    a.set_label("PALDATA", base + reglen + patlen + ntlen)
    return a.resolve()


def build_rom(patterns: bytes, nametable: bytes, palette: bytes) -> bytes:
    code = _code(len(VDP_REGS), len(patterns), len(nametable))
    rom = bytearray(ROM_SIZE)
    blob = code + bytes(VDP_REGS) + bytes(patterns) + bytes(nametable) + bytes(palette)
    if len(blob) > 0x7FF0:
        raise ValueError("SMS image data overruns the 32KB cartridge")
    rom[0:len(blob)] = blob

    # "TMR SEGA" header at $7FF0 (export region, 32KB) so the BIOS accepts it.
    rom[0x7FF0:0x7FF8] = b"TMR SEGA"
    rom[0x7FF8:0x7FFA] = bytes(2)                    # reserved
    chk = sum(rom[0:0x7FF0]) & 0xFFFF                # checksum of $0000-$7FEF
    rom[0x7FFA] = chk & 0xFF
    rom[0x7FFB] = (chk >> 8) & 0xFF
    rom[0x7FFC:0x7FFF] = bytes(3)                    # product code / version
    rom[0x7FFF] = 0x4C                               # region 4 (export) | size $C (32K)
    return bytes(rom)


def write_sms(rom: bytes, path: str) -> str:
    with open(path, "wb") as f:
        f.write(rom)
    return path