"""Assemble the Atari 6502 boot viewers with `xa` (origin $2000, no load prefix)."""

from __future__ import annotations

import os
import shutil
import subprocess
import tempfile

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

SOURCES = {
    "gr15": "gr15.s",
    "gr9": "gr9.s",
    "gr8": "gr8.s",
    "gr15dli": "gr15dli.s",
}

_cache: dict[tuple, bytes] = {}

# How long the viewer holds the picture (see atari/viewer/awyt.i).
WAIT_MODES = {"forever": 0, "key": 1, "seconds": 2}

# Slideshow advance behaviour and per-mode multi-image viewer parameters
# (source, ANTIC mode byte, GPRIOR, colour-register layout).
SS_WAITMODE = {"key": 1, "seconds": 2, "both": 3}
SLIDESHOW_PARAMS = {
    "gr15": (0x0E, 0x00, 0),
    "gr9":  (0x0F, 0x40, 1),
    "gr8":  (0x0F, 0x00, 2),
}
SLIDESHOW_SOURCES = dict.fromkeys(SLIDESHOW_PARAMS, "slideshow_static.s")
SLIDESHOW_SOURCES["gr15dli"] = "slideshow_dli.s"   # DLI mode, its own engine


class AssemblerError(RuntimeError):
    pass


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


def assemble_stub(viewer_key: str, display: str = "forever", seconds: int = 0,
                  video: str = "ntsc") -> bytes:
    waitmode = WAIT_MODES.get(display, 0)
    rate = 50 if video == "pal" else 60
    key = (viewer_key, waitmode, int(seconds), rate)
    if key in _cache:
        return _cache[key]
    if not have_xa():
        raise AssemblerError(
            "The 'xa' (xa65) assembler was not found on PATH.\n"
            "Install it with:  sudo apt install xa65")
    if not os.path.exists(os.path.join(VIEWER_DIR, SOURCES[viewer_key])):
        raise AssemblerError(f"viewer source missing: {SOURCES[viewer_key]}")

    # Wrapper sets the options then includes the real source; runs from VIEWER_DIR
    # so the source's #include "awyt.i" resolves (xa looks relative to cwd).
    wrapper = (
        f"#define WAITMODE {waitmode}\n"
        f"#define WAITSECS {max(0, int(seconds))}\n"
        f"#define RATE {rate}\n"
        f'#include "{SOURCES[viewer_key]}"\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 for {viewer_key}:\n{proc.stdout}{proc.stderr}")
            with open(out, "rb") as f:
                raw = f.read()
        finally:
            os.unlink(wrap)
    _cache[key] = raw
    return raw


def _xa(wrapper: str, what: str) -> bytes:
    """Assemble a generated wrapper with xa (run from VIEWER_DIR so #includes
    resolve); return raw bytes."""
    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 for {what}:\n{proc.stdout}{proc.stderr}")
            with open(out, "rb") as f:
                return f.read()
        finally:
            os.unlink(wrap)


def build_slideshow_stub(viewer_key: str, n_images: int, base_sec: int, spi: int,
                         advance: str = "both", seconds: int = 10,
                         loop: bool = True, video: str = "ntsc") -> bytes:
    """Assemble the multi-image slideshow viewer (origin $2000, no load prefix).

    ``base_sec`` is the disk sector of image 0 and ``spi`` the sectors per image
    (both fixed by the ATR layout); the viewer SIO-reads image i from
    base_sec + i*spi.  ``advance``/``seconds`` set the per-slide dwell, ``loop``
    whether it wraps.
    """
    if viewer_key not in SLIDESHOW_SOURCES:
        raise AssemblerError(f"no Atari slideshow viewer for mode {viewer_key}")
    rate = 50 if video == "pal" else 60
    common = (f"#define WAITMODE {SS_WAITMODE[advance]}\n"
              f"#define WAITSECS {max(0, int(seconds))}\n"
              f"#define RATE {rate}\n"
              f"#define NIMAGES {n_images}\n"
              f"#define LOOPFLAG {1 if loop else 0}\n"
              f"#define BASESEC {base_sec}\n"
              f"#define SPI {spi}\n")
    if viewer_key in SLIDESHOW_PARAMS:             # static gr15/gr9/gr8
        dlmode, gprior, colormode = SLIDESHOW_PARAMS[viewer_key]
        common += (f"#define DLMODE ${dlmode:02X}\n"
                   f"#define GPRIOR ${gprior:02X}\n"
                   f"#define COLORMODE {colormode}\n")
    return _xa(common + f'#include "{SLIDESHOW_SOURCES[viewer_key]}"\n',
               f"slideshow_{viewer_key}")


CART_SIZE = 0x4000               # 16K Atari cartridge ROM at $8000-$BFFF


def build_cart_rom(viewer_key: str, data: bytes, display: str = "forever",
                   seconds: int = 0, video: str = "ntsc") -> bytes:
    """Assemble the loader + the disk viewer stub + picture data into a 16K
    cartridge ROM with the Atari run/init footer at $BFFA."""
    stub = assemble_stub(viewer_key, display, seconds, video)
    wrapper = (
        f"#define STUB_PAGES {(len(stub) + 255) // 256}\n"
        f"#define DATA_PAGES {(len(data) + 255) // 256}\n"
        f"#define STUB_LEN {len(stub)}\n"
        '#include "cart.s"\n')
    loader = _xa(wrapper, "cart")
    rom = bytearray(loader + stub + bytes(data))
    if len(rom) > CART_SIZE:
        raise AssemblerError(
            f"viewer + image = {len(rom)} bytes, over the 16K cartridge limit")
    rom += bytes(CART_SIZE - len(rom))
    # Atari cartridge footer at $BFFA (ROM offset $3FFA).
    rom[0x3FFA] = 0x00; rom[0x3FFB] = 0x80     # CARTCS run address = $8000
    rom[0x3FFC] = 0x00                          # cart present
    rom[0x3FFD] = 0x04                          # option byte: start the cartridge
    rom[0x3FFE] = 0x00; rom[0x3FFF] = 0x80     # CARTAD init address = $8000
    return bytes(rom)