#!/usr/bin/env python3
"""Build the JumPHP assets of Sorcery from the openSorcery resources folder.

Source: the C++/SDL2 port of Sorcery+ (Virgin Games 1985, Amstrad CPC),
https .../david/sorcery (its `resources/` folder and `assets/audio`).
Everything is derived from that folder; the original .dsk is NOT needed.

    python3 convert.py /mnt/openspace1/david/sorcery

Outputs into ../.. (the game folder):
  assets/tiles/tiles1.png tiles2.png        128x128 RGBA, 16x16 tiles of 8x8 (pen -> scene palette)
  assets/sprites/sprites1.png sprites2.png  288x192 RGBA, 12x8 frames of 24x24 (pen 0 transparent)
  assets/sprites/extra.png                  288x48 RGBA, frames 96.. of the sprite list
  assets/fonts/font1.png font2.png font3.png  128x72 RGBA atlases of 8x8 glyphs (text palette, pens mask 1/2/3)
  assets/pictures/title.png                 320x200 RGBA (the mode-1 picture rebuilt from the 160x200 sheet)
  assets/data/CODE.bin SMESSC.bin TISCRN.bin halloffame.bin   copied as is (parsed at run time)
  assets/rooms/roomNN.txt                   copied as is
  assets/levels/<pack>/...                  copied, tiles.png converted like the tiles sheets
  assets/sfx/*.wav assets/music/title.mp3   copied from assets/audio
"""
import os, shutil, sys
from PIL import Image

SRC = sys.argv[1] if len(sys.argv) > 1 else "/mnt/openspace1/david/sorcery"
DST = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
CODE_BASE = 0x7000

def fail(msg):
    print("convert: " + msg)
    sys.exit(1)

# ---------------------------------------------------------------- colour
def lv_of(i):
    return [0, 128, 255][i]

def ink_rgb(ink):
    """Firmware ink 0..26 -> (r, g, b). Port of Assets::inkColour (C++ writes ABGR; red is the low byte)."""
    ink %= 27
    r = lv_of((ink // 3) % 3)
    g = lv_of(ink // 9)
    b = lv_of(ink % 3)
    return (r, g, b)

INK2HW = [0x14, 0x04, 0x15, 0x1C, 0x18, 0x1D, 0x0C, 0x05, 0x0D, 0x16, 0x06, 0x17, 0x1E, 0x00,
          0x1F, 0x0E, 0x07, 0x0F, 0x12, 0x02, 0x13, 0x1A, 0x19, 0x1B, 0x0A, 0x03, 0x0B]

def hardware_rgb(hw):
    """Hardware colour number -> (r, g, b). Port of Assets::hardwareColour."""
    hw &= 0x1F
    if hw == 0x01: hw = 0x00
    elif hw == 0x08: hw = 0x05
    elif hw == 0x09: hw = 0x03
    elif hw == 0x10: hw = 0x04
    elif hw == 0x11: hw = 0x02
    for i in range(27):
        if INK2HW[i] == hw:
            return ink_rgb(i)
    return (0, 0, 0)

# ---------------------------------------------------------------- CODE.bin
class Code:
    def __init__(self, path):
        with open(path, "rb") as f:
            self.d = f.read()

    def at(self, a):
        off = a - CODE_BASE
        return self.d[off] if 0 <= off < len(self.d) else 0

    def strings(self, a, count):
        out = []
        for _ in range(count):
            s = ""
            while True:
                b = self.at(a); a += 1
                s += chr(b & 0x7F)
                if b & 0x80:
                    break
            out.append(s)
        return out

# ---------------------------------------------------------------- indexed PNGs
def load_pen_png(path):
    """Indexed PNG -> list of palette indices + size. The PNG palette order IS the pen order."""
    im = Image.open(path)
    if im.mode == "P":
        return im.load(), im.size
    if im.mode == "RGB" or im.mode == "RGBA":
        # map every pixel to the nearest colour of its own palette is not possible here:
        # the sources are always indexed; converting through RGB would need the preview palette
        pal = im.getpalette()
        rgb = im.convert("RGBA")
        return ("rgb", rgb.load(), pal, im.size)
    fail(path + ": expected an indexed PNG, got " + im.mode)

def pen_at(source, x, y):
    if source[0] == "rgb":
        _, px, pal, size = source
        if x < 0 or y < 0 or x >= size[0] or y >= size[1]:
            return 0
        r, g, b = px[x, y][:3]
        best, bi = 1 << 30, 0
        for i in range(0, len(pal) // 3 if pal else 16):
            dr, dg, db = pal[i * 3] - r, pal[i * 3 + 1] - g, pal[i * 3 + 2] - b
            d = dr * dr + dg * dg + db * db
            if d < best:
                best, bi = d, i
        return bi
    px, size = source
    if x < 0 or y < 0 or x >= size[0] or y >= size[1]:
        return 0
    return px[x, y]

def make_dir(p):
    os.makedirs(p, exist_ok=True)

# ---------------------------------------------------------------- main
def main():
    res = os.path.join(SRC, "resources")
    audio = os.path.join(SRC, "assets", "audio")
    for p in (res, audio):
        if not os.path.isdir(p):
            fail("missing " + p)

    code = Code(os.path.join(res, "data", "CODE.bin"))

    scene_pal = [ink_rgb(code.at(0x92F8 + i)) for i in range(16)]
    text_pal = [hardware_rgb(code.at(0x9370 + i) & 0x1F) for i in range(4)]
    picture_pal = [ink_rgb(code.at(0x930C + i)) for i in range(4)]
    border = hardware_rgb(code.at(0x936F) & 0x1F)

    def emit_png(path, w, h, get):
        im = Image.new("RGBA", (w, h))
        px = im.load()
        for y in range(h):
            for x in range(w):
                px[x, y] = get(x, y)
        im.save(path)

    print("== graphics")
    gdir = os.path.join(DST, "assets")
    for sub in ("tiles", "sprites", "fonts", "pictures", "data", "rooms", "sfx", "music"):
        make_dir(os.path.join(gdir, sub))

    for set_n in (1, 2):
        # tiles: 64x128 pens (16x16 grid of 4x8) -> 128x128 screen pixels (each pen 2 px wide, 8 lines tall)
        src = load_pen_png(os.path.join(res, "graphics", "tiles%d.png" % set_n))
        out = os.path.join(gdir, "tiles", "tiles%d.png" % set_n)
        emit_png(out, 128, 128, lambda x, y, s=src, p=scene_pal: p[pen_at(s, x // 2, y)] + (255,))
        # sprites: 144x192 pens (12x8 frames of 12x24) -> 288x192 (pen 0 transparent)
        src = load_pen_png(os.path.join(res, "graphics", "sprites%d.png" % set_n))
        out = os.path.join(gdir, "sprites", "sprites%d.png" % set_n)

        def sprite_pix(x, y, s=src, p=scene_pal):
            pen = pen_at(s, x // 2, y)
            return p[pen] + (255,) if pen else (0, 0, 0, 0)
        emit_png(out, 288, 192, sprite_pix)

    # extra sprites (frames 96..): 144x48 pens -> 288x48
    src = load_pen_png(os.path.join(res, "graphics", "extra.png"))
    emit_png(os.path.join(gdir, "sprites", "extra.png"), 288, 48,
             lambda x, y: (lambda pen: scene_pal[pen] + (255,) if pen else (0, 0, 0, 0))(pen_at(src, x // 2, y)))

    # font: three atlases, one per pens mask (the colour of a pixel is text_pal[g & pens])
    src = load_pen_png(os.path.join(res, "graphics", "font.png"))
    for pens in (1, 2, 3):
        emit_png(os.path.join(gdir, "fonts", "font%d.png" % pens), 128, 72,
                 lambda x, y, m=pens: (lambda g: text_pal[g & m] + (255,) if g & m else (0, 0, 0, 0))(pen_at(src, x, y)))

    # title picture: the 160x200 sheet shows each mode-1 screen byte as two mode-0 pixels; rebuild the bytes, read as mode 1
    src = load_pen_png(os.path.join(res, "graphics", "title.png"))
    pic = [[0] * 320 for _ in range(200)]
    for y in range(200):
        for bx in range(80):
            p0 = pen_at(src, bx * 2, y)
            p1 = pen_at(src, bx * 2 + 1, y)
            b = ((p0 & 1) << 7 | (p0 >> 1 & 1) << 3 | (p0 >> 2 & 1) << 5 | (p0 >> 3 & 1) << 1 |
                 (p1 & 1) << 6 | (p1 >> 1 & 1) << 2 | (p1 >> 2 & 1) << 4 | (p1 >> 3 & 1))
            for p in range(4):
                pic[y][bx * 4 + p] = ((b >> (3 - p)) & 1) | ((b >> (7 - p)) & 1) << 1
    emit_png(os.path.join(gdir, "pictures", "title.png"), 320, 200,
             lambda x, y: picture_pal[pic[y][x]] + (255,))

    print("== data")
    for name in ("CODE.bin", "SMESSC.bin", "TISCRN.bin", "halloffame.bin"):
        shutil.copyfile(os.path.join(res, "data", name), os.path.join(gdir, "data", name))

    print("== rooms")
    for n in range(79):
        shutil.copyfile(os.path.join(res, "rooms", "room%02d.txt" % n), os.path.join(gdir, "rooms", "room%02d.txt" % n))

    print("== levels")
    packs = [l.strip() for l in open(os.path.join(res, "levels", "index.txt")) if l.strip() and not l.strip().startswith("#")]
    for pack in packs:
        base = os.path.join(res, "levels", pack)
        dst = os.path.join(gdir, "levels", pack)
        make_dir(os.path.join(dst, "rooms"))
        shutil.copyfile(os.path.join(base, "level.txt"), os.path.join(dst, "level.txt"))
        if os.path.isfile(os.path.join(base, "extras.txt")):
            shutil.copyfile(os.path.join(base, "extras.txt"), os.path.join(dst, "extras.txt"))
        n = 0
        while os.path.isfile(os.path.join(base, "rooms", "room%02d.txt" % n)):
            shutil.copyfile(os.path.join(base, "rooms", "room%02d.txt" % n), os.path.join(dst, "rooms", "room%02d.txt" % n))
            n += 1
        tiles_png = os.path.join(base, "tiles.png")
        if os.path.isfile(tiles_png):
            src = load_pen_png(tiles_png)
            emit_png(os.path.join(dst, "tiles.png"), 128, 128,
                     lambda x, y, s=src: scene_pal[pen_at(s, x // 2, y)] + (255,))
        print("   pack %-8s %2d rooms" % (pack, n))

    print("== audio")
    sfx = {"shoot": "shoot", "rise": "rise", "click": "click", "enemy_explode": "enemy_explode",
           "player_death": "player_death", "pickup": "pickup", "blip": "blip", "burst": "burst",
           "hit": "hit", "zap": "zap", "thud": "thud", "counter": "counter", "game_over": "game_over",
           "level_done": "level_done", "whoosh": "whoosh"}
    for name in sfx:
        shutil.copyfile(os.path.join(audio, "sfx_%s.wav" % name), os.path.join(gdir, "sfx", "%s.wav" % name))
    shutil.copyfile(os.path.join(audio, "music_title.mp3"), os.path.join(gdir, "music", "title.mp3"))

    print("scene palette:   " + " ".join("%02d=%s" % (i, "".join("%02x" % c for c in scene_pal[i])) for i in (0, 3, 7, 15)))
    print("text palette:    %s border %s" % (text_pal, border))
    print("done -> " + DST)

if __name__ == "__main__":
    main()