#!/usr/bin/env python3
"""
split_sd81.py — ZX81+SD81 Booster binary splitter + loader generator

Splits the flat binary generated by zxbc.py (--arch zx81sd, -f bin)
into 8KB pages for loading from the ZX81 BASIC loader, and generates
the BASIC listing (plain text, not tokenized) needed to load them.

Page map:
  Page 8  -> block 0 ($0000-$1FFF): vectors + stage2 + runtime + user code
  Page 9  -> block 1 ($2000-$3FFF): continuation if the binary exceeds 8KB
  Page 10 -> block 2 ($4000-$5FFF): same
  Page 11 -> block 3 ($6000-$7FFF): same (external stage 1 resides here before the jump)
  Page 12 -> block 4 ($8000-$9FFF): sysvars + heap (data, not executable without MC45)
  Page 13 -> block 5 ($A000-$BFFF): heap continuation

The binary starts at $0000 and includes no .tap/.tzx header. Each
output file is named <BASE>P<N>.BIN (uppercase), where N is the page.

The generated BASIC loader (plain text, not yet tokenized to .p)
follows the sequence manually validated on hardware/emulator:

  2 FAST
  5 LOAD THEN CLEAR 24575
  10 LOAD FAST "BOOT1.BIN"CODE 24576
  20 LOAD *MAP 7,8
  25 LOAD FAST "<BASE>P8.BIN"CODE 57344
  30 LOAD *MAP 7,9
  35 LOAD FAST "<BASE>P9.BIN"CODE 57344
  ...
  <N0> LOAD *MAP 7,63
  <N0+10> RAND USR 24576

Note: "LOAD THEN CLEAR", "LOAD *MAP" and "LOAD FAST ... CODE" are
extensions of the SD81 Booster firmware over the ZX81's BASIC (they
don't exist in the original ROM). Block 7 ($E000/57344) is used as a
temporary paging window to dump each physical page before BOOT1.BIN
does the final mapping of blocks 0-5 through port $E7. The final
"LOAD *MAP 7,63" forces the switch to full-paging mode (64 pages) and
leaves block 7 in a neutral state before jumping to the program.

Besides the plain-text listing, this also generates the real tokenized
.p file (ready to load/run) using zx81_p_loader.py, a Python port of
EightyOne's tokenizer (zx81BasicLoader.cpp / IBasicLoader.cpp).

To load and run the resulting package on the ZX81 (or in EightyOne
pointing at the SD image): `LOAD FAST "<PREFIX>"` — nothing else needs
selecting afterwards, it runs straight from there.

Usage:
  python split_sd81.py <input.bin> [output_base]

  If output_base isn't given, the input's name without extension is
  used.

Example:
  python split_sd81.py test_sd81.bin
  -> TEST_SD81P8.BIN  (first page, always present)
  -> TEST_SD81P9.BIN  (only if the binary exceeds 8 KB)
  -> ...
  -> test_sd81_loader.txt  (loader's BASIC listing, plain text)
  -> TEST_SD81.P           (the same listing, tokenized)
"""

import os
import sys

from zx81_p_loader import build_p_file

PAGE_SIZE = 8192  # 8KB per page
FIRST_PAGE = 8  # SD81 page assigned to block 0
CLEANUP_PAGE = 63  # neutral page: forces full-paging mode (64 pages)
BOOT1_FILE = "BOOT1.BIN"
BOOT1_ADDR = 24576  # $6000
CLEAR_ADDR = 24575  # protects BOOT1.BIN (reserved right below it)
LOAD_WINDOW_ADDR = 57344  # $E000, block 7's window


def split(input_path: str, output_base: str) -> list[str]:
    with open(input_path, "rb") as f:
        data = f.read()

    if not data:
        sys.exit(f"Error: {input_path} is empty")

    pages_written = []
    page_num = FIRST_PAGE
    offset = 0

    while offset < len(data):
        chunk = data[offset : offset + PAGE_SIZE]

        # Pad up to PAGE_SIZE with 0xFF (undefined value for unwritten FLASH/RAM)
        if len(chunk) < PAGE_SIZE:
            chunk = chunk + b"\xff" * (PAGE_SIZE - len(chunk))

        out_path = f"{output_base}P{page_num}.BIN"
        with open(out_path, "wb") as f:
            f.write(chunk)

        pages_written.append(out_path)
        print(
            f"  Page {page_num} (block {page_num - FIRST_PAGE}): "
            f"{out_path}  "
            f"[{offset:#06x} - {min(offset + PAGE_SIZE - 1, len(data) - 1):#06x}]"
        )

        offset += PAGE_SIZE
        page_num += 1

    return pages_written


def generate_loader_lines(page_files: list[str]) -> list[tuple[int, str]]:
    lines = []
    lines.append((2, "FAST"))
    lines.append((5, f"LOAD THEN CLEAR {CLEAR_ADDR}"))
    lines.append((10, f'LOAD FAST "{BOOT1_FILE}"CODE {BOOT1_ADDR}'))

    line_no = 20
    for i, page_file in enumerate(page_files):
        page_num = FIRST_PAGE + i
        lines.append((line_no, f"LOAD *MAP 7,{page_num}"))
        lines.append((line_no + 5, f'LOAD FAST "{page_file}"CODE {LOAD_WINDOW_ADDR}'))
        line_no += 10

    lines.append((line_no, f"LOAD *MAP 7,{CLEANUP_PAGE}"))
    lines.append((line_no + 10, f"RAND USR {BOOT1_ADDR}"))

    return lines


def generate_loader_text(lines: list[tuple[int, str]]) -> str:
    return "\n".join(f"{n} {cmd}" for n, cmd in lines) + "\n"


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)

    input_path = sys.argv[1]
    output_base = sys.argv[2] if len(sys.argv) > 2 else os.path.splitext(input_path)[0]
    output_base = os.path.basename(output_base).upper()

    # Page and loader.p file names are written in the ZX81 charset
    # (letters, digits and a few symbols: no underscore or other
    # non-representable characters). Fail early and clearly instead of
    # blowing up later while tokenizing the loader.
    from zx81_p_loader import ascii_to_zx

    for c in output_base:
        try:
            ascii_to_zx(c)
        except ValueError:
            sys.exit(
                f"Error: base name {output_base!r} contains {c!r}, "
                "not representable in the ZX81 charset (avoid '_' and similar)."
            )

    if not os.path.isfile(input_path):
        sys.exit(f"Error: {input_path!r} not found")

    size = os.path.getsize(input_path)
    num_pages = (size + PAGE_SIZE - 1) // PAGE_SIZE
    print(f"Binary: {input_path}  ({size} bytes, {size / 1024:.1f} KB, {num_pages} page(s))")
    print(f"Splitting into {PAGE_SIZE}-byte pages (starting SD81 page = {FIRST_PAGE})...")
    print()

    pages = split(input_path, output_base)

    print()
    print(f"Generated {len(pages)} file(s).")
    if len(pages) == 1:
        print("The program fits in a single page (8 KB).")
    else:
        print(f"Pages {FIRST_PAGE} to {FIRST_PAGE + len(pages) - 1}.")

    loader_lines = generate_loader_lines(pages)
    loader_text = generate_loader_text(loader_lines)

    # The .p is a real file name for the ZX81 (charset with no
    # underscore): derived from output_base (same prefix as the
    # pages), concatenated with no separator, just like "<BASE>P8.BIN".
    base_no_ext = os.path.splitext(input_path)[0]
    loader_txt_path = f"{base_no_ext}_loader.txt"
    loader_p_path = f"{output_base}.P"

    with open(loader_txt_path, "w", newline="\n") as f:
        f.write(loader_text)

    p_data = build_p_file(loader_lines)
    with open(loader_p_path, "wb") as f:
        f.write(p_data)

    print()
    print(f"BASIC loader (plain text): {loader_txt_path}")
    print("---")
    print(loader_text, end="")
    print("---")
    print(f"Tokenized BASIC loader (.p): {loader_p_path}  ({len(p_data)} bytes)")
    print()
    print("NOTE: 'LOAD THEN CLEAR', 'LOAD *MAP' and 'LOAD FAST ... CODE' are")
    print("extensions of the SD81 Booster firmware, encoded with the same token")
    print("mechanism as standard ZX81 BASIC (see zx81_p_loader.py).")
    print(f"Copy {BOOT1_FILE}, the page files and {os.path.basename(loader_p_path)}")
    print("to the SD card.")


if __name__ == "__main__":
    main()
