#!/usr/bin/env python3
# extractAssets.py - pull human-readable assets out of the disk: Dan Bunten's letter (track 1),
# the hidden EA message (track 1 sector 17) and the loader title picture (as PNG).
import os, struct, zlib
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
EXT  = os.path.join(ROOT, "extracted")
DOCS = os.path.join(ROOT, "docs")
os.makedirs(DOCS, exist_ok=True)

def raw(t, s):
    return open(f"{EXT}/sectors/t{t:02d}s{s:02d}.bin", "rb").read()

# --- letter: track 1 sectors 1-10, plain ASCII, CR = newline, ends before the padding in sector 10
text = b"".join(raw(1, s) for s in range(1, 11))
end = text.find(b"72203") + 5
letter = text[:end].decode("latin-1").replace("\r", "\n")
with open(f"{DOCS}/track1_DanBuntenLetter.txt", "w") as f:
    f.write("Hidden text stored in plain ASCII on track 1, sectors 1-10 of the Modem Wars disk\n")
    f.write("(never shown by the game; addressed to anyone who cracks or inspects the disk).\n")
    f.write("=" * 78 + "\n\n")
    f.write(letter + "\n")

# --- EA message in the drive bootstrap sector
d = raw(1, 17)
msg = d[0x4D:].split(b"\x00")[0].decode("latin-1")
with open(f"{DOCS}/track1_sector17_message.txt", "w") as f:
    f.write("Text embedded after the code in the 1541 bootstrap sector (track 1 sector 17):\n\n")
    f.write(msg + "\n")

# --- title picture: multicolour bitmap at LOAD $A000, screen RAM $9C00, colour RAM $9800, background $D021=0
load = open(f"{EXT}/files/load.prg", "rb").read()[2:]
bitmap = load[0xA000-0x9800:0xA000-0x9800+8000]
screen = load[0x9C00-0x9800:0x9C00-0x9800+1000]
color  = load[0x9800-0x9800:0x9800-0x9800+1000]
PALETTE = [(0,0,0),(255,255,255),(136,0,0),(170,255,238),(204,68,204),(0,204,85),(0,0,170),(238,238,119),
           (221,136,85),(102,68,0),(255,119,119),(51,51,51),(119,119,119),(170,255,102),(0,136,255),(187,187,187)]
W, H = 320, 200
pixels = bytearray(W * H * 3)
for cy in range(25):
    for cx in range(40):
        cell = cy * 40 + cx
        c1 = screen[cell] >> 4
        c2 = screen[cell] & 15
        c3 = color[cell] & 15
        for row in range(8):
            b = bitmap[cell * 8 + row]
            for px in range(4):
                bits = (b >> (6 - px * 2)) & 3
                col = [0, c1, c2, c3][bits]
                r, g, bl = PALETTE[col]
                for dup in range(2):
                    x = cx * 8 + px * 2 + dup
                    y = cy * 8 + row
                    i = (y * W + x) * 3
                    pixels[i:i+3] = bytes((r, g, bl))
def png(path, w, h, rgb):
    rows = b"".join(b"\x00" + bytes(rgb[y*w*3:(y+1)*w*3]) for y in range(h))
    def chunk(tag, data):
        c = struct.pack(">I", len(data)) + tag + data
        return c + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF)
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        f.write(chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)))
        f.write(chunk(b"IDAT", zlib.compress(rows, 9)))
        f.write(chunk(b"IEND", b""))
# scale 2x vertically too so the aspect is right
scaled = bytearray()
for y in range(H):
    row = pixels[y*W*3:(y+1)*W*3]
    scaled += row + row
png(f"{DOCS}/titleScreen.png", W, H * 2, scaled)
print("wrote", os.listdir(DOCS))