"""Write a bootable Atari ``.atr`` disk image natively (no external tools).

A self-booting Atari disk needs no DOS: sector 1 begins with a 6-byte boot header
(flags, sector-count, load-address, init-address); the OS loads ``count`` 128-byte
sectors to the load address and JSRs ``load+6``.  We pack the whole viewer+picture
blob that way, so inserting the disk shows the picture.
"""

from __future__ import annotations

SECTOR_SIZE = 128
TOTAL_SECTORS = 720           # single density, ~90K
LOAD_ADDR = 0x2000            # boot load address (blob origin)
DATA_ADDR = 0x4000            # where the bitmap must land (4K-aligned for ANTIC)


class AtrError(RuntimeError):
    pass


def build_blob(stub: bytes, data: bytes) -> bytes:
    """Combine the assembled viewer ``stub`` (origin $2000, starts with the 6-byte
    boot header) + padding + picture ``data`` (which must reside from $4000)."""
    pad = (DATA_ADDR - LOAD_ADDR) - len(stub)
    if pad < 0:
        raise AtrError(f"viewer stub {len(stub)} bytes exceeds "
                       f"{DATA_ADDR - LOAD_ADDR} before $4000")
    return stub + bytes(pad) + bytes(data)


def _write_atr_sectors(path: str, data: bytes) -> str:
    """Write ``data`` (already laid out as raw sectors) as a single-density ATR,
    padded to the full 720-sector disk with the 16-byte ATR header."""
    data = bytearray(data)
    if len(data) > TOTAL_SECTORS * SECTOR_SIZE:
        raise AtrError("slideshow exceeds the 720-sector disk capacity")
    data += bytes(TOTAL_SECTORS * SECTOR_SIZE - len(data))
    paragraphs = (TOTAL_SECTORS * SECTOR_SIZE) // 16
    header = bytes([
        0x96, 0x02,
        paragraphs & 0xFF, (paragraphs >> 8) & 0xFF,
        SECTOR_SIZE & 0xFF, (SECTOR_SIZE >> 8) & 0xFF,
        (paragraphs >> 16) & 0xFF, (paragraphs >> 24) & 0xFF,
        0, 0, 0, 0, 0, 0, 0, 0,
    ])
    with open(path, "wb") as f:
        f.write(header)
        f.write(data)
    return path


def write_slideshow_atr(path: str, stub: bytes, images: list[bytes],
                        boot_sectors: int, spi: int) -> str:
    """Write a self-booting slideshow ATR.

    Sectors 1..boot_sectors hold ``stub`` (boot header byte 1 patched to
    boot_sectors so the OS loads it all); each image then occupies ``spi``
    consecutive 128-byte sectors (image i at sector boot_sectors + 1 + i*spi),
    matching what the viewer SIO-reads.
    """
    if len(stub) > boot_sectors * SECTOR_SIZE:
        raise AtrError(f"slideshow viewer {len(stub)} bytes exceeds "
                       f"{boot_sectors} boot sectors")
    blob = bytearray(stub)
    blob[1] = boot_sectors                            # OS loads this many sectors
    blob += bytes(boot_sectors * SECTOR_SIZE - len(blob))
    for img in images:
        if len(img) > spi * SECTOR_SIZE:
            raise AtrError("image larger than its sector allotment")
        blob += bytes(img) + bytes(spi * SECTOR_SIZE - len(img))
    return _write_atr_sectors(path, bytes(blob))


def write_boot_atr(path: str, blob: bytes) -> str:
    """Write ``blob`` as a bootable single-density ATR. Patches the boot sector
    count (byte 1) from the blob length."""
    nsec = (len(blob) + SECTOR_SIZE - 1) // SECTOR_SIZE
    if nsec > 255:
        raise AtrError(f"boot blob needs {nsec} sectors (max 255)")
    if nsec > TOTAL_SECTORS:
        raise AtrError("boot blob exceeds disk capacity")

    blob = bytearray(blob)
    blob[1] = nsec                                   # boot header sector count
    blob += bytes((-len(blob)) % SECTOR_SIZE)        # pad to whole sectors

    data = bytearray(blob)
    data += bytes(TOTAL_SECTORS * SECTOR_SIZE - len(data))   # pad disk to 720 sectors

    paragraphs = (TOTAL_SECTORS * SECTOR_SIZE) // 16
    header = bytes([
        0x96, 0x02,                                  # magic
        paragraphs & 0xFF, (paragraphs >> 8) & 0xFF,  # size (paragraphs) low
        SECTOR_SIZE & 0xFF, (SECTOR_SIZE >> 8) & 0xFF,
        (paragraphs >> 16) & 0xFF, (paragraphs >> 24) & 0xFF,
        0, 0, 0, 0, 0, 0, 0, 0,
    ])
    with open(path, "wb") as f:
        f.write(header)
        f.write(data)
    return path