"""Assemble the PET viewer with `xa` and build the loadable .prg.

The PRG loads at the PET BASIC start ($0401): a BASIC stub ``10 SYS1056`` then
the 6502 viewer (at $0420 = 1056) then the screen data (at $0500).  Running it
(or SYS 1056) copies the screen codes to $8000.
"""
from __future__ import annotations

import os
import shutil
import subprocess
import tempfile

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

BASIC_START = 0x0401
ML_ORG = 0x0420                # 1056
DATA_ORG = 0x0500

# BASIC: 10 SYS1056  (bytes from $0401)
_STUB = bytes([0x0B, 0x04, 0x0A, 0x00, 0x9E,
               0x31, 0x30, 0x35, 0x36, 0x00, 0x00, 0x00])


class AssemblerError(RuntimeError):
    pass


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


def _assemble(pages: int) -> bytes:
    if not have_xa():
        raise AssemblerError("The 'xa' (xa65) assembler was not found (apt install xa65).")
    wrapper = (f"#define DATA ${DATA_ORG:04X}\n"
               f"#define PAGES {pages}\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_prg(screen: bytes) -> bytes:
    """screen = screen-RAM bytes (1000 for 40-col, 2000 for 80-col)."""
    pages = (len(screen) + 255) // 256
    data = bytes(screen) + bytes(pages * 256 - len(screen))     # pad to whole pages
    code = _assemble(pages)
    mem = bytearray()
    mem += _STUB
    mem += b"\x00" * (ML_ORG - BASIC_START - len(mem))
    mem += code
    if len(mem) > DATA_ORG - BASIC_START:
        raise AssemblerError("viewer code overruns the data area")
    mem += b"\x00" * (DATA_ORG - BASIC_START - len(mem))
    mem += data
    return bytes([BASIC_START & 0xFF, BASIC_START >> 8]) + bytes(mem)