"""Assemble the Atari 2600 kernel with `xa` and lay out the cartridge.

A static image is a 4K cart (one kernel + nine 192-byte tables).  An interlace
image is an 8K F8-bankswitch cart: two 4K banks, each holding the SAME kernel
plus one of the two table sets (frame A / frame B).  The kernel toggles the bank
once per frame, so the two frames alternate at 60Hz and blend to ~4-6 perceived
colours per scanline.
"""

from __future__ import annotations

import os
import shutil
import subprocess
import tempfile

VIEWER_DIR = os.path.dirname(os.path.abspath(__file__))
WAIT_MODES = {"forever": 0, "fire": 1, "key": 1, "seconds": 2}

# page-aligned data tables ($F100..$F900)
TABLES = {"PF0L": 0xF100, "PF1L": 0xF200, "PF2L": 0xF300, "PF0R": 0xF400,
          "PF1R": 0xF500, "PF2R": 0xF600, "BKL": 0xF700, "BKR": 0xF800,
          "PFT": 0xF900}
ORDER = ["PF0L", "PF1L", "PF2L", "PF0R", "PF1R", "PF2R", "BKL", "BKR", "PFT"]


class AssemblerError(RuntimeError):
    pass


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


def _assemble_kernel(display: str, seconds: int, interlace: bool) -> 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)
    defs = "".join(f"#define {n} ${a:04X}\n" for n, a in TABLES.items())
    wrapper = (defs +
               f"#define WAITMODE {waitmode}\n"
               f"#define WAITSECS {max(0, int(seconds))}\n"
               f"#define IL {1 if interlace else 0}\n"
               '#include "a2600.s"\n')
    with tempfile.TemporaryDirectory() as td:
        out = os.path.join(td, "k.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:
                kernel = f.read()
        finally:
            os.unlink(wrap)
    if len(kernel) > 0x100:
        raise AssemblerError(f"kernel is {len(kernel)} bytes, overruns the table "
                             "page at $F100")
    return kernel


def _lay_bank(kernel: bytes, data: bytes) -> bytearray:
    """One 4096-byte bank: kernel at $F000, nine tables, reset/IRQ vectors."""
    rom = bytearray(b"\x00" * 4096)
    rom[0:len(kernel)] = kernel
    tables = [data[i * 192:(i + 1) * 192] for i in range(9)]
    for name, tab in zip(ORDER, tables):
        off = TABLES[name] - 0xF000
        rom[off:off + 192] = tab
    rom[0xFFC] = 0x00          # reset vector lo -> $F000
    rom[0xFFD] = 0xF0
    rom[0xFFE] = 0x00          # IRQ/BRK vector
    rom[0xFFF] = 0xF0
    return rom


def build_cart(data: bytes, display: str = "forever", seconds: int = 0) -> bytes:
    """data = nine 192-byte tables (static, 4K cart) or eighteen (two sets ->
    interlace 8K F8 cart).  Returns the cart image with reset vectors set."""
    if len(data) == 9 * 192:
        kernel = _assemble_kernel(display, seconds, interlace=False)
        return bytes(_lay_bank(kernel, data))
    if len(data) == 18 * 192:
        kernel = _assemble_kernel(display, seconds, interlace=True)
        bank0 = _lay_bank(kernel, data[:9 * 192])
        bank1 = _lay_bank(kernel, data[9 * 192:])
        return bytes(bank0 + bank1)
    raise ValueError(f"expected 1728 or 3456 bytes of tables, got {len(data)}")