"""Assemble the Atari 5200 viewer with `xa` and lay out the 32K cartridge.

The cartridge fills $4000-$BFFF.  The 5200 BIOS jumps to the address at $BFFE-F
(duplicated at $BFE8-9, with $BFFC=0 / $BFFD=$FF) -- verified in MAME a5200.

ANTIC DMAs the bitmap and display list straight from cartridge ROM, so we only
place them at fixed cart addresses and point ANTIC at them:
    $4000  viewer code
    $4100  GTIA register script (offset, value pairs, $FF-terminated)
    $4200  display list  (ANTIC reads it here)
    $6000  bitmap        (the converter's 4K-split blob -> $6000 / $7000)
"""

from __future__ import annotations

import os
import shutil
import subprocess
import tempfile

VIEWER_DIR = os.path.dirname(os.path.abspath(__file__))

CART_BASE = 0x4000
CART_SIZE = 0x8000
SCRIPT_ADDR = 0x4100
DLIST_ADDR = 0x4200
BITMAP_ADDR = 0x6000          # 4K-aligned: blob $4000/$5000 -> $6000/$7000
SPLIT_LINE = 102              # matches atari _common.split_screen
LINES = 192

# ANTIC mode byte + GTIA register script per display mode.  Script entries are
# (GTIA offset from $C000, colour value); colours come from the converter blob.
ANTIC_MODE = {"gr15": 0x0E, "gr8": 0x0F, "gr9": 0x0F}

WAIT_MODES = {"forever": 0, "key": 1, "seconds": 2}


def _gtia_script(mode: str, colors) -> bytes:
    c = list(colors)
    if mode == "gr15":              # 4 colours: COLBK, COLPF0, COLPF1, COLPF2
        regs = [(0x1A, c[0]), (0x16, c[1]), (0x17, c[2]), (0x18, c[3]), (0x1B, 0x00)]
    elif mode == "gr8":             # 2 colours: background (COLPF2+COLBK), fg (COLPF1)
        regs = [(0x18, c[0]), (0x1A, c[0]), (0x17, c[1]), (0x1B, 0x00)]
    else:                           # gr9: hue in COLBK, GTIA mode 9 via PRIOR
        regs = [(0x1A, c[0]), (0x1B, 0x40)]
    out = bytearray()
    for off, val in regs:
        out += bytes([off & 0xFF, val & 0xFF])
    out.append(0xFF)                # terminator
    return bytes(out)


def _dlist(mode: str) -> bytes:
    m = ANTIC_MODE[mode]
    dl = bytearray([0x70, 0x70, 0x70])                      # 24 blank scan lines
    dl += bytes([0x40 | m, 0x00, 0x60])                     # LMS -> $6000
    dl += bytes([m] * (SPLIT_LINE - 1))                     # 102 lines from $6000
    dl += bytes([0x40 | m, 0x00, 0x70])                     # LMS -> $7000
    dl += bytes([m] * (LINES - SPLIT_LINE - 1))             # 90 lines from $7000
    dl += bytes([0x41, DLIST_ADDR & 0xFF, DLIST_ADDR >> 8])  # JVB -> dlist (loop)
    return bytes(dl)


class AssemblerError(RuntimeError):
    pass


def have_xa() -> bool:
    return shutil.which("xa") is not None


def _assemble(display: str = "forever", seconds: int = 0,
              video: str = "ntsc") -> bytes:
    if not have_xa():
        raise AssemblerError("The 'xa' (xa65) assembler was not found on PATH.\n"
                             "Install it with:  sudo apt install xa65")
    waitmode = WAIT_MODES.get(display, 0)
    rate = 50 if video == "pal" else 60
    wrapper = (f"#define SCRIPT ${SCRIPT_ADDR:04X}\n"
               f"#define DLIST ${DLIST_ADDR:04X}\n"
               f"#define WAITMODE {waitmode}\n"
               f"#define WAITSECS {max(1, int(seconds))}\n"
               f"#define RATE {rate}\n"
               '#include "viewer.s"\n')
    with tempfile.TemporaryDirectory() as td:
        out = os.path.join(td, "v.bin")
        fd, wrap = tempfile.mkstemp(suffix=".s", prefix="_wrap_", dir=VIEWER_DIR)
        try:
            with os.fdopen(fd, "w") as f:
                f.write(wrapper)
            proc = subprocess.run(["xa", "-o", out, os.path.basename(wrap)],
                                  capture_output=True, text=True, cwd=VIEWER_DIR)
            if proc.returncode != 0:
                raise AssemblerError(f"xa failed:\n{proc.stdout}{proc.stderr}")
            with open(out, "rb") as f:
                return f.read()
        finally:
            os.unlink(wrap)


def build_cart(mode: str, data: bytes, display: str = "forever",
               seconds: int = 0, video: str = "ntsc") -> bytes:
    """data = the atari converter blob (bitmap 4K-split at offset 0/0x1000, colour
    register bytes from offset 0x2000)."""
    if mode not in ANTIC_MODE:
        raise ValueError(f"unsupported 5200 mode {mode}")
    code = _assemble(display, seconds, video)
    bitmap = data[:0x2000]
    colors = data[0x2000:]
    script = _gtia_script(mode, colors)
    dlist = _dlist(mode)

    rom = bytearray(b"\xff" * CART_SIZE)
    def place(addr, blob):
        off = addr - CART_BASE
        rom[off:off + len(blob)] = blob
    place(CART_BASE, code)
    place(SCRIPT_ADDR, script)
    place(DLIST_ADDR, dlist)
    place(BITMAP_ADDR, bitmap)

    # cartridge header (verified in MAME a5200)
    def put16(addr, val):
        off = addr - CART_BASE
        rom[off] = val & 0xFF
        rom[off + 1] = (val >> 8) & 0xFF
    put16(0xBFFE, CART_BASE)        # start address -- BIOS jumps here
    put16(0xBFE8, CART_BASE)        # duplicate (validation)
    rom[0xBFFC - CART_BASE] = 0x00
    rom[0xBFFD - CART_BASE] = 0xFF
    return bytes(rom)