#!/usr/bin/env python3
"""
pop_pack_kid.py — упаковщик спрайтов персонажа Kid (Фаза K0 порта).

Аналог pop_pack_bg.py, но для chtab_2_kid (база ресурса 400): image id N
= res(400+N).png (data/KID), палитра res400.pal (16 цветов) -> слоты
Sprinter 0x70-0x7F (env 0x50 / wall 0x60 заняты фоном).  Пиксель i:
0->0xFF (прозрач), i->0x70+i.

Прямая адресация (ноль remap в W2): kid[img>>5], idx img&31.  Дырки —
фейковые записи (как в фоне).  Переиспользует pack_atlas/blob/scale6to8
из pop_pack_bg.

Выход: poc/res/kid/kid0..N.atl + kid.pal + kid_atlas.h
"""
import os
import sys
from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import pop_pack_bg as B   # pack_atlas, blob, scale6to8, ENV_SHIFT

KID_DIR = "/Users/alex/Projects/DIY/Z80/Sprinter/C-Compiler/applications/PoP/SDLPoP/data/KID"
OUT_DIR = os.path.join(HERE, "..", "poc", "res", "kid")
PAL_BASE = 0x70          # слоты палитры Sprinter под Kid
SWORD_PAL_BASE = 0x80    # chtab_0 (меч в руке) -> слоты 0x80..0x8F
GUARD_PAL_BASE = 0x90    # chtab_5 (страж) -> слоты 0x90..0x9F
SWORD_IDS = tuple(range(20, 28))  # sword_tbl для кадров 229..236 (найти-убрать)
SHIFT = 3                # kid[img>>3], idx img&7 (Kid-кадры крупные: 8/стр ≤16КБ)
MASK = (1 << SHIFT) - 1


def transpose_cols(w, h, rowmajor):
    """row-major (pix[row*w+col]) -> column-major (out[col*h+row]).
    Спрайты Kid хранятся КОЛОНКАМИ: gfx_blit_cols рисует вертикальным accel-
    проходом (даёт бесплатный горизонтальный флип реверсом порядка колонок;
    см. memory/accel_vertical_copy).  Заголовок getimage (w,h) не меняется."""
    out = bytearray(w * h)
    for col in range(w):
        base = col * h
        for row in range(h):
            out[base + row] = rowmajor[row * w + col]
    return bytes(out)


def load_kid(aidx):
    """Атлас-индекс aidx (== frame.image) -> res(401+aidx).png -> (w, h, 8bpp
    COLUMN-MAJOR).  0->0xFF, i->0x70+i.  Пиксели транспонированы в колонки.

    МАППИНГ (seg008.c:1589): персонаж рисуется add_midtable(chtab, obj_id+1),
    внутри get_image(chtab, id-1)=get_image(chtab, obj_id) -> chtab-индекс =
    frame.image; первая картинка chtab_2_kid = res401 (res400 = палитра).
    Значит frame.image=N -> res(401+N).  Раньше грузили res(400+N) — off-by-one:
    стойка frame_15 (image 14) рисовалась res414 (беговой наклон) вместо res415
    (прямая стойка)."""
    path = os.path.join(KID_DIR, f"res{401 + aidx}.png")
    if not os.path.exists(path):
        return None
    im = Image.open(path)
    if im.mode != "P":
        B.die(f"res{401+aidx}: mode {im.mode}, ожидался P")
    src = im.tobytes()
    rowmajor = bytes((0xFF if v == 0 else PAL_BASE + v) for v in src)
    return im.width, im.height, transpose_cols(im.width, im.height, rowmajor)


def load_sword(id):
    """chtab_0 (меч в руке): sword_tbl[].id = N -> res(700+N+1).png в data/PRINCE
    (в оригинале add_midtable(chtab, obj_id+1) -> get_image(chtab, obj_id)).
    Пиксели, как у Kid, COLUMN-MAJOR (gfx_blit_cols + бесплатный флип)."""
    path = os.path.join(os.path.dirname(B.R.VDUNGEON_DIR), B.POT_DIR_NAME,
                        f"res{700 + id + 1}.png")
    if not os.path.exists(path):
        B.die(f"нет файла {path}")
    im = Image.open(path)
    if im.mode != "P":
        B.die(f"res{700+id+1}: mode {im.mode}, ожидался P")
    src = im.tobytes()
    rowmajor = bytes((0xFF if v == 0 else SWORD_PAL_BASE + v) for v in src)
    return im.width, im.height, transpose_cols(im.width, im.height, rowmajor)


def _read16(path):
    b = open(path, "rb").read()
    return [(b[4 + i * 3], b[4 + i * 3 + 1], b[4 + i * 3 + 2]) for i in range(16)]


def build_palette(out_path):
    """ЕДИНАЯ игровая палитра (gfx_pal_fload пишет с индекса 0, перезатирая
    всё — bg и Kid обязаны жить в одном файле): chtab_1@0x40 (пламя/зелья),
    env@0x50, wall@0x60, kid@0x70.
    """
    pal = bytearray(256 * 4)
    groups = [
        (0x50, os.path.join(B.R.VDUNGEON_DIR, "res200.pal")),   # env
        (0x60, os.path.join(B.R.VDUNGEON_DIR, "res360.pal")),   # wall
        (PAL_BASE, os.path.join(KID_DIR, "res400.pal")),        # kid 0x70
        (B.POT_PAL_BASE, os.path.join(os.path.dirname(B.R.VDUNGEON_DIR),
                                      B.POT_DIR_NAME, "res150.pal")),  # chtab_1 0x40
        (SWORD_PAL_BASE, os.path.join(os.path.dirname(B.R.VDUNGEON_DIR),
                                      B.POT_DIR_NAME, "res700.pal")),  # chtab_0 0x80
    ]
    # 16 базовых цветов экрана (mono-блит красит ими: пламя, пузырьки зелья)
    for i, (r, g, bl) in enumerate(B.VGA16):
        o = (B.VGA16_PAL_BASE + i) * 4
        pal[o:o + 4] = bytes((B.scale6to8(bl), B.scale6to8(g), B.scale6to8(r), 0))
    for base, path in groups:
        for i, (r, g, bl) in enumerate(_read16(path)):
            o = (base + i) * 4
            pal[o:o + 4] = bytes((B.scale6to8(bl), B.scale6to8(g), B.scale6to8(r), 0))
    # Палитра стража берётся НЕ из .pal, а из res10.bin (guard_palettes:
    # 7 палитр по 16 цветов, 6-бит) по level.guards_color; на уровне 1 у
    # обоих стражей color = 2.  См. toolchain/pop_pack_guard.py.
    import pop_pack_guard as G
    for i, (r, g, bl) in enumerate(G.guard_palette()):
        o = (GUARD_PAL_BASE + i) * 4
        pal[o:o + 4] = bytes((B.scale6to8(bl), B.scale6to8(g), B.scale6to8(r), 0))

    with open(out_path, "wb") as f:
        f.write(pal)


def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    # атлас-индексы aidx (== frame.image), для которых есть res(401+aidx).png
    ids = []
    for aidx in range(0, 255):
        if os.path.exists(os.path.join(KID_DIR, f"res{401 + aidx}.png")):
            ids.append(aidx)
    if not ids:
        B.die("нет спрайтов KID")

    pages = {}
    for id in ids:
        w, h, pix = load_kid(id)
        pages.setdefault(id >> SHIFT, {})[id & MASK] = (w, h, pix)

    max_pg = max(pages)
    names = []
    total = 0
    for pg in range(max_pg + 1):
        if pg in pages:
            nm = f"kid{pg}.atl"
            cnt, sz = B.pack_atlas(os.path.join(OUT_DIR, nm), pages[pg])
            total += sz
            names.append(nm)
            print(f"  {nm}: {len(pages[pg])} спрайтов, каталог {cnt}, {sz} Б")
        else:
            names.append(None)
            print(f"  (kid стр {pg}: пусто)")

    # chtab_0: меч в руке.  Пакуем ТОЛЬКО кадры подъёма/ножен (sword_tbl
    # 35..42 -> id 20..27); индекс в атласе = id - 20 (бой не портирован).
    sw = {}
    for i, id in enumerate(SWORD_IDS):
        w, h, pix = load_sword(id)
        sw[i] = (w, h, pix)
    cnt, sz = B.pack_atlas(os.path.join(OUT_DIR, "sword.atl"), sw)
    total += sz
    print(f"  sword.atl (chtab_0, id {SWORD_IDS[0]}..{SWORD_IDS[-1]}): "
          f"{len(sw)} спрайтов, каталог {cnt}, {sz} Б")

    build_palette(os.path.join(OUT_DIR, "kid.pal"))

    # C-заголовок
    hdr = os.path.join(OUT_DIR, "kid_atlas.h")
    with open(hdr, "w") as f:
        f.write("/* kid_atlas.h — раскладка атласов Kid. Сгенерировано "
                "pop_pack_kid.py. */\n#ifndef KID_ATLAS_H\n#define KID_ATLAS_H\n")
        f.write(f"#define KID_SHIFT   {SHIFT}\n#define KID_MASK    {MASK}\n")
        f.write(f"#define KID_PAGES   {len(names)}\n#define KID_PAL     0x{PAL_BASE:02X}\n")
        f.write("static const char *const kid_atl[KID_PAGES] = {\n")
        for nm in names:
            f.write("    %s,\n" % ('"%s"' % nm if nm else "0"))
        f.write('};\n#define KID_PAL_FILE "kid.pal"\n')
        f.write('#define KID_SWORD_ATL "sword.atl"   /* chtab_0: меч в руке */\n')
        f.write(f"#define KID_SWORD_ID0 {SWORD_IDS[0]}  /* индекс в атласе = sword_tbl.id - ID0 */\n")
        f.write("#endif\n")

    print(f"ИТОГО: {len(ids)} спрайтов Kid, {len(names)} EMM-страниц, "
          f"{total} Б диск (+1КБ палитра) -> {OUT_DIR}")


if __name__ == "__main__":
    main()