#!/usr/bin/env python3
# assetbake.py -- bake indexed PNG assets into JoeyLib native binary
# formats. Replaces the .jas pipeline (png2jas + the runtime asset
# loader's chunky->planar conversion path).
#
# Each target's runtime can use the baked bytes directly:
#   * Tile banks  (.tbk) are per-target planar (or chunky for DOS/IIgs).
#     The Phase 11 walker reads sprite tile data as chunky 4bpp, but
#     tile-paste is platform-private -- Amiga is plane-major, ST is
#     row-major-with-planes-per-row, DOS / IIgs are chunky 4bpp.
#   * Sprite cels (.spr) are cross-target chunky 4bpp. The sprite walker
#     reads chunky and does c2p inline at draw time, so the same blob
#     works on every platform.
#
# Both formats embed an optional 16-entry $0RGB palette so the runtime
# can pin tile / sprite colors without depending on the stage palette
# state.
#
# Source PNGs must be 'P' (indexed) mode with at most 16 colors and
# pixel dimensions that are multiples of 8 (tile size). Tile banks
# extract one 8x8 tile per grid cell row-major. Sprite sheets extract
# one cel per grid cell, where cel size is given via --cell WxH (in
# tile units, not pixels).
#
# --------------------------------------------------------------------
# .tbk format
# --------------------------------------------------------------------
#   offset  bytes  field
#   0       4      magic "JTB1"
#   4       1      target           (1=amiga, 2=atarist, 3=dos, 4=iigs)
#   5       1      hasPalette       (0 or 1)
#   6       2      tileCount        (LE16)
#   8       4      reserved         (zero)
#   12      32     palette[16] LE16 $0RGB (zeros if hasPalette==0)
#   44      ...    tileCount * 32 bytes, layout per target:
#                    amiga  : plane-major  [p0_row0..7][p1_row0..7]...
#                    atarist: row-major-planes-per-row
#                             [r0_p0 r0_p1 r0_p2 r0_p3][r1_p0..]
#                    dos    : chunky 4bpp  [r0_b0..3][r1_b0..3]...
#                    iigs   : chunky 4bpp  (same as dos)
#
# --------------------------------------------------------------------
# .spr format
# --------------------------------------------------------------------
#   offset  bytes  field
#   0       4      magic "JSP1"
#   4       1      target           (0 = cross-platform chunky 4bpp)
#   5       1      hasPalette
#   6       1      widthTiles       (cel width in 8-px tiles)
#   7       1      heightTiles      (cel height in 8-px tiles)
#   8       2      cellCount        (LE16)
#   10      2      reserved         (zero)
#   12      32     palette[16] LE16 $0RGB (zeros if hasPalette==0)
#   44      ...    cellCount * (widthTiles * heightTiles * 32) bytes
#                  per cel: tile-major (tile(0,0), tile(1,0), ...,
#                  tile(W-1,0), tile(0,1), ...). Per tile: 8 rows of
#                  4 chunky 4bpp bytes, high nibble = left pixel.

import argparse
import struct
import sys

from PIL import Image


MAGIC_TBK             = b"JTB1"
MAGIC_SPR             = b"JSP1"
TARGET_CROSS_CHUNKY   = 0
TARGET_AMIGA          = 1
TARGET_ATARIST        = 2
TARGET_DOS            = 3
TARGET_IIGS           = 4

TARGET_NAMES = {
    "amiga":   TARGET_AMIGA,
    "atarist": TARGET_ATARIST,
    "dos":     TARGET_DOS,
    "iigs":    TARGET_IIGS,
}

TILE_PX        = 8         # 8x8 tile
TILE_BYTES     = 32        # 4bpp packed = 4 bytes/row * 8 rows
PALETTE_BYTES  = 32        # 16 entries * 2 bytes LE
HEADER_BYTES   = 44        # both formats share header length


# ----- PNG helpers ----------------------------------------------------

def loadIndexedPng(path):
    img = Image.open(path)
    if img.mode != "P":
        sys.exit(f"assetbake: {path} mode={img.mode}, need 'P' (indexed)")
    w, h = img.size
    if w <= 0 or h <= 0:
        sys.exit(f"assetbake: {path} zero-sized")
    if (w % TILE_PX) != 0 or (h % TILE_PX) != 0:
        sys.exit(f"assetbake: {path} dims {w}x{h} not 8-px-tile multiples")
    px = img.load()
    grid = [[px[x, y] for x in range(w)] for y in range(h)]
    return grid, w, h, img.getpalette()


def packPaletteRgb444(rawPalette):
    # PNG palette is RGB triples (256 entries). Pack first 16 into
    # $0RGB little-endian 16-bit words (4 bits per channel).
    out = bytearray()
    if rawPalette is None:
        return bytes(PALETTE_BYTES)
    for i in range(16):
        base = i * 3
        if base + 2 < len(rawPalette):
            r = rawPalette[base    ] >> 4
            g = rawPalette[base + 1] >> 4
            b = rawPalette[base + 2] >> 4
        else:
            r = g = b = 0
        word = (r << 8) | (g << 4) | b
        out += struct.pack("<H", word)
    return bytes(out)


# ----- Tile encoders --------------------------------------------------

def encodeTileChunky(grid, ox, oy):
    # 8 rows of 4 bytes, high nibble = left pixel.
    out = bytearray()
    for r in range(TILE_PX):
        row = grid[oy + r]
        for c in range(0, TILE_PX, 2):
            hi = row[ox + c    ] & 0x0F
            lo = row[ox + c + 1] & 0x0F
            out.append((hi << 4) | lo)
    return bytes(out)


def encodeTilePlaneMajor(grid, ox, oy):
    # Amiga: 4 planes * 8 rows, one byte per row per plane.
    # Plane 0 bytes 0..7, plane 1 bytes 8..15, etc.
    out = bytearray()
    for k in range(4):
        for r in range(TILE_PX):
            row = grid[oy + r]
            b = 0
            for c in range(TILE_PX):
                if (row[ox + c] >> k) & 1:
                    b |= (1 << (7 - c))
            out.append(b)
    return bytes(out)


def encodeTileRowPlanesPerRow(grid, ox, oy):
    # ST: 8 rows, each row is 4 plane bytes in order [p0 p1 p2 p3].
    out = bytearray()
    for r in range(TILE_PX):
        row = grid[oy + r]
        for k in range(4):
            b = 0
            for c in range(TILE_PX):
                if (row[ox + c] >> k) & 1:
                    b |= (1 << (7 - c))
            out.append(b)
    return bytes(out)


TILE_ENCODERS = {
    TARGET_AMIGA:   encodeTilePlaneMajor,
    TARGET_ATARIST: encodeTileRowPlanesPerRow,
    TARGET_DOS:     encodeTileChunky,
    TARGET_IIGS:    encodeTileChunky,
}


# ----- Bakers ---------------------------------------------------------

def bakeTileBank(srcPath, dstPath, target):
    grid, w, h, rawPal = loadIndexedPng(srcPath)
    tilesAcross = w // TILE_PX
    tilesDown   = h // TILE_PX
    tileCount   = tilesAcross * tilesDown
    if tileCount > 0xFFFF:
        sys.exit(f"assetbake: {srcPath} {tileCount} tiles exceeds 65535")
    encoder = TILE_ENCODERS[target]
    palette = packPaletteRgb444(rawPal)
    hasPal  = 1 if rawPal is not None else 0

    header = bytearray()
    header += MAGIC_TBK
    header += struct.pack("<BBHI", target, hasPal, tileCount, 0)
    header += palette
    assert len(header) == HEADER_BYTES

    body = bytearray()
    for ty in range(tilesDown):
        for tx in range(tilesAcross):
            body += encoder(grid, tx * TILE_PX, ty * TILE_PX)

    with open(dstPath, "wb") as f:
        f.write(header)
        f.write(body)


def bakeSpriteSheet(srcPath, dstPath, cellWTiles, cellHTiles):
    grid, w, h, rawPal = loadIndexedPng(srcPath)
    cellWPx = cellWTiles * TILE_PX
    cellHPx = cellHTiles * TILE_PX
    if (w % cellWPx) != 0 or (h % cellHPx) != 0:
        sys.exit(f"assetbake: {srcPath} {w}x{h} not multiple of "
                 f"cell {cellWPx}x{cellHPx}")
    cellsAcross = w // cellWPx
    cellsDown   = h // cellHPx
    cellCount   = cellsAcross * cellsDown
    if cellCount > 0xFFFF:
        sys.exit(f"assetbake: {srcPath} {cellCount} cels exceeds 65535")

    palette = packPaletteRgb444(rawPal)
    hasPal  = 1 if rawPal is not None else 0

    header = bytearray()
    header += MAGIC_SPR
    header += struct.pack("<BBBB", TARGET_CROSS_CHUNKY, hasPal,
                          cellWTiles, cellHTiles)
    header += struct.pack("<HH", cellCount, 0)
    header += palette
    assert len(header) == HEADER_BYTES

    body = bytearray()
    # Cels in PNG reading order: left-to-right, top-to-bottom.
    # Each cel = tile-major chunky 4bpp.
    for cy in range(cellsDown):
        for cx in range(cellsAcross):
            celOx = cx * cellWPx
            celOy = cy * cellHPx
            for ty in range(cellHTiles):
                for tx in range(cellWTiles):
                    body += encodeTileChunky(grid,
                                             celOx + tx * TILE_PX,
                                             celOy + ty * TILE_PX)

    with open(dstPath, "wb") as f:
        f.write(header)
        f.write(body)


# ----- CLI ------------------------------------------------------------

def parseCell(spec):
    # "WxH" -> (W, H) in tile units.
    parts = spec.lower().split("x")
    if len(parts) != 2:
        sys.exit(f"assetbake: bad --cell '{spec}', need 'WxH'")
    try:
        w = int(parts[0])
        h = int(parts[1])
    except ValueError:
        sys.exit(f"assetbake: bad --cell '{spec}', non-integer")
    if w <= 0 or h <= 0:
        sys.exit(f"assetbake: bad --cell '{spec}', zero/negative")
    return w, h


def main():
    p = argparse.ArgumentParser(description="Bake JoeyLib native assets.")
    p.add_argument("--type", required=True, choices=["sprite", "tile"])
    p.add_argument("--target", choices=list(TARGET_NAMES.keys()),
                   help="required for --type tile")
    p.add_argument("--cell", help="cel size as WxH in tile units (sprite only)")
    p.add_argument("input")
    p.add_argument("output")
    args = p.parse_args()

    if args.type == "tile":
        if args.target is None:
            sys.exit("assetbake: --type tile requires --target")
        bakeTileBank(args.input, args.output, TARGET_NAMES[args.target])
    else:
        if args.cell is None:
            sys.exit("assetbake: --type sprite requires --cell WxH")
        wT, hT = parseCell(args.cell)
        bakeSpriteSheet(args.input, args.output, wT, hT)


if __name__ == "__main__":
    main()