#!/usr/bin/env python3
"""Упаковка story/PV intro SDLPoP в полосы SPA1 для Sprinter.

Исходники остаются единственным источником истины: TITLE/res41..45, res51 и
PV/res800/850/950/980. Runtime получает только готовые 320x200 композиции
и 256-цветные палитры. Это намеренно переносит смешивание палитров Princess,
Jaffar, спальни и кровати из ограниченного Z80 в build-time.
"""
from pathlib import Path

from PIL import Image

ROOT = Path(__file__).resolve().parent.parent
TITLE = ROOT / "SDLPoP" / "data" / "TITLE"
PV = ROOT / "SDLPoP" / "data" / "PV"
PRINCE = ROOT / "SDLPoP" / "data" / "PRINCE"
KID = ROOT / "SDLPoP" / "data" / "KID"
OUT = ROOT / "poc" / "res" / "pv"

ATL_MAGIC = b"SPA1"
ATL_DIR_OFF = 0x68
ATL_MIN_SIZE = 0x100
PAGE_SIZE = 16384
STRIP_H = 49

# ФОН ТЕКСТОВОЙ РАМКИ — ОТДЕЛЬНЫЙ ИНДЕКС.  У оригинала это индекс 14 в
# строке палитры спрайтов title40, и load_title_images() красит его либо в
# #100060 (интро), либо в #800000 (финал, seg001:586).  Раньше мы
# ремапили 14 в 9 — «тёмно-синий» палитры res51; подменить его нельзя,
# потому что тем же индексом нарисована сама титульная картинка (8265
# пикселей).  Поэтому фон уезжает в свободный индекс, и финал красит
# только его (pop_pal_story_ending).
STORY_BG_INDEX = 16
STORY_BG_COLOR = (0x10, 0x00, 0x60)
# Цвет текста таблицы рекордов: 0xB7 оригинала = седьмой цвет палитры
# title (224,160,0).  Глифы шрифта запечены индексом POP_FONT_COLOR, так
# что цвет им даёт именно эта запись; тёмный набор (POP_FONT_DARK_COLOR)
# остаётся чёрным — старшие записи story.pal и так чёрные.
STORY_FONT_INDEX = 0x3F


def transpose_cols(width: int, height: int, rowmajor: bytes) -> bytes:
    """Row-major -> column-major для бесплатного horizontal flip в BGI."""
    out = bytearray(width * height)
    for column in range(width):
        base = column * height
        for row in range(height):
            out[base + row] = rowmajor[row * width + column]
    return bytes(out)


def atlas_blob(pixels: bytes, width: int, height: int) -> bytes:
    """Один непрозрачный getimage в минимальном SPA1-атласе."""
    data_off = ATL_DIR_OFF + 8
    out = bytearray(data_off)
    out[:4] = ATL_MAGIC
    out[4] = 1
    out[ATL_DIR_OFF:ATL_DIR_OFF + 2] = data_off.to_bytes(2, "little")
    out[ATL_DIR_OFF + 2] = min(width, 255)
    out[ATL_DIR_OFF + 3] = height
    out[ATL_DIR_OFF + 4] = out[ATL_DIR_OFF + 5] = 1
    out += width.to_bytes(2, "little") + height.to_bytes(2, "little")
    out += pixels
    if len(out) < ATL_MIN_SIZE:
        out += bytes(ATL_MIN_SIZE - len(out))
    if len(out) > PAGE_SIZE:
        raise SystemExit(f"intro atlas {len(out)} B > {PAGE_SIZE} B")
    return bytes(out)


def sprite_atlas(path: Path, entries: dict[int, tuple[int, int, bytes]]) -> None:
    """Атлас небольших прозрачных кадров, целиком живущий в одной EMM.

    В отличие от atlas_blob() для полноэкранной полосы, здесь нужен каталог
    нескольких картинок: runtime держит страницу открытой всю PV-сцену и
    переключает только idx.  Пустые индексы не требуются — все наборы ниже
    нумеруются подряд с нуля.
    """
    count = len(entries)
    if not count or count > 255:
        raise SystemExit(f"{path}: bad sprite atlas count {count}")
    data_off = ATL_DIR_OFF + count * 8
    directory = bytearray(count * 8)
    blobs = bytearray()
    for idx in range(count):
        width, height, pixels = entries[idx]
        offset = data_off + len(blobs)
        blobs += width.to_bytes(2, "little") + height.to_bytes(2, "little")
        blobs += pixels
        d = idx * 8
        directory[d:d + 2] = offset.to_bytes(2, "little")
        directory[d + 2] = width
        directory[d + 3] = height
        directory[d + 4] = directory[d + 5] = 1
    out = bytearray(data_off)
    out[:4] = ATL_MAGIC
    out[4] = count
    out[ATL_DIR_OFF:ATL_DIR_OFF + len(directory)] = directory
    out += blobs
    if len(out) < ATL_MIN_SIZE:
        out += bytes(ATL_MIN_SIZE - len(out))
    if len(out) > PAGE_SIZE:
        raise SystemExit(f"{path}: {len(out)} B > one EMM page")
    path.write_bytes(out)


def sprite_frames(directory: Path, first: int, last: int,
                  color_base: int, column_major: bool = False
                  ) -> dict[int, tuple[int, int, bytes]]:
    """PNG-кадры -> getimage. Ноль PNG становится прозрачным 0xFF BGI."""
    result = {}
    for index, number in enumerate(range(first, last + 1)):
        image = png(directory, number)
        pixels = bytes(0xFF if pixel == 0 else color_base + pixel
                       for pixel in image.tobytes())
        if column_major:
            pixels = transpose_cols(image.width, image.height, pixels)
        result[index] = image.width, image.height, pixels
    return result


def png(directory: Path, number: int) -> Image.Image:
    image = Image.open(directory / f"res{number}.png")
    if image.mode != "P":
        raise SystemExit(f"res{number}.png: expected indexed PNG, got {image.mode}")
    return image.copy()


def pal_file(path: Path) -> list[tuple[int, int, int]]:
    """16 цветов из .pal ресурса (4 байта заголовка, дальше RGB по 6 бит)."""
    raw = path.read_bytes()
    return [(raw[4 + i * 3] << 2, raw[5 + i * 3] << 2, raw[6 + i * 3] << 2)
            for i in range(16)]


def palette(image: Image.Image) -> list[tuple[int, int, int]]:
    raw = image.getpalette()
    return [tuple(raw[index * 3:index * 3 + 3]) for index in range(16)]


def write_palette(path: Path, entries: list[tuple[int, int, int]],
                  extra=None) -> None:
    """B,G,R,0 — ровно формат gfx_pal_fload; незадействованное чёрное."""
    out = bytearray()
    entries = (entries + [(0, 0, 0)] * 256)[:256]
    for index, color in (extra or {}).items():
        entries[index] = color
    for red, green, blue in entries:
        out += bytes((blue, green, red, 0))
    path.write_bytes(out)


def write_scene(prefix: str, number: int, canvas: Image.Image) -> None:
    if canvas.size != (320, 200):
        raise SystemExit(f"{prefix}{number}: expected 320x200, got {canvas.size}")
    for part, top in enumerate(range(0, 200, STRIP_H)):
        band = canvas.crop((0, top, 320, min(top + STRIP_H, 200)))
        (OUT / f"{prefix}{number}_{part}.atl").write_bytes(
            atlas_blob(band.tobytes(), band.width, band.height))


def story_scenes() -> list[Image.Image]:
    """show_title/end_sequence в единой с title палитре.

    res41 кодирует тёмный фон индексом 14, а его палитра делает этот индекс
    чёрным; show_title() SDLPoP подменяет его #100060 либо #800000. Тексты
    res42..45 имеют индекс 1, но этот индекс у res41 тоже тёмный; ремап на
    индекс title/res51 15 (белый) позволяет проявлять историю прямо поверх
    последнего title-экрана без смены палитры. Фон уезжает в
    STORY_BG_INDEX, чтобы финал мог покрасить его отдельно.

    Последний экран — фон таблицы рекордов: та же рамка плюс логотип
    PRINCE OF PERSIA на y=24 (HOF_POP оригинала — тот же спрайт res54, что
    и в титрах, только выше).
    """
    frame_raw = bytearray(png(TITLE, 41).tobytes())
    for offset, pixel in enumerate(frame_raw):
        if pixel == 14:
            frame_raw[offset] = STORY_BG_INDEX
    frame = Image.frombytes("P", (320, 200), bytes(frame_raw))
    result = []
    for text_number, ypos in ((42, 25), (43, 25), (45, 26), (44, 25)):
        canvas = frame.copy()
        raw = bytearray(canvas.tobytes())
        text = png(TITLE, text_number)
        for sy in range(text.height):
            for sx in range(text.width):
                if text.getpixel((sx, sy)):
                    raw[(ypos + sy) * 320 + 24 + sx] = 15
        canvas.frombytes(bytes(raw))
        result.append(canvas)
    result.append(frame)

    hof = frame.copy()
    raw = bytearray(hof.tobytes())
    logo = png(TITLE, 54)
    for sy in range(logo.height):
        for sx in range(logo.width):
            pixel = logo.getpixel((sx, sy))
            if pixel:                       # blitters_10h_transp: 0 сквозной
                raw[(24 + sy) * 320 + 24 + sx] = pixel
    hof.frombytes(bytes(raw))
    result.append(hof)
    return result


def pv_palette_canvas() -> tuple[Image.Image, list[tuple[int, int, int]]]:
    """Спальня + bed. Номера 0..15=room, 16..31=Princess,
    32..47=Jaffar, 48..63=bed. Благодаря этому у смешанных PNG не теряются
    исходные цвета при одном gfx_pal_fload() на весь PV."""
    room = png(PV, 951)
    bed = png(PV, 981)
    raw = bytearray(room.tobytes())
    bed_raw = bed.tobytes()
    for sy in range(bed.height):
        for sx in range(bed.width):
            pixel = bed_raw[sy * bed.width + sx]
            if pixel:
                raw[(142 + sy) * 320 + sx] = 48 + pixel
    canvas = Image.frombytes("P", room.size, bytes(raw))
    # 80..95 — пламя факелов. Это отдельный chtab PRINCE в оригинале; без
    # отдельной строки цвета динамические кадры становились бы цветами стены.
    # 112..127 (0x70) — палитра КИДА: на PV-экране финала и сцен 8/9 он и
    # мышь рисуются СВОИМИ атласами (chtab_2), а их индексы запечены под
    # игровые слоты 0x70.  Без этой строки они выходили бы чёрными.
    return (canvas, palette(room) + palette(png(PV, 801)) +
            palette(png(PV, 851)) + palette(bed) + palette(png(PV, 901)) +
            palette(png(PRINCE, 151)) + [(0, 0, 0)] * 16 +
            pal_file(KID / "res400.pal"))


def sprite_paste(dst: Image.Image, sprite: Image.Image, x: int, floor: int,
                 color_base: int) -> None:
    """Положить спрайт ногами на floor, remap его палитру в отдельную строку."""
    raw = bytearray(dst.tobytes())
    src = sprite.tobytes()
    y = floor - sprite.height
    for sy in range(sprite.height):
        dy = y + sy
        if dy < 0 or dy >= 200:
            continue
        for sx in range(sprite.width):
            dx = x + sx
            pixel = src[sy * sprite.width + sx]
            if pixel and 0 <= dx < 320:
                raw[dy * 320 + dx] = color_base + pixel
    dst.frombytes(bytes(raw))


def pv_scenes() -> tuple[list[Image.Image], list[tuple[int, int, int]]]:
    base, colors = pv_palette_canvas()

    # Позиции Char из init_princess()/init_vizier(): x=120/198, y=166.
    # В оригинале точный offset задаёт frame table; статические композиции
    # выравнивают персонажей по полу y=166, не хранят таблицы PV-анимации.
    waiting = base.copy()
    sprite_paste(waiting, png(PV, 801), 110, 166, 16)

    jaffar = waiting.copy()
    sprite_paste(jaffar, png(PV, 851), 188, 166, 32)

    magic = base.copy()
    sprite_paste(magic, png(PV, 805), 110, 166, 16)
    sprite_paste(magic, png(PV, 870), 188, 166, 32)
    # res953 — первый полный кадр песочных часов; он использует палитру room.
    sprite_paste(magic, png(PV, 953), 18, 166, 0)

    alone = base.copy()
    sprite_paste(alone, png(PV, 816), 110, 166, 16)
    sprite_paste(alone, png(PV, 959), 18, 166, 0)

    # FG9: статические ключевые позы тех же scene 2/4/6/8/9/12. PV2
    # использует второй набор Princess (res900), поэтому он лежит в строке
    # палитры 64..79. Мышь в оригинале меняет только короткую позу рядом с
    # Princess; здесь соответствующая стадия выделена отдельным кадром.
    stand = base.copy()
    sprite_paste(stand, png(PV, 801), 134, 166, 16)

    lying = base.copy()
    sprite_paste(lying, png(PV, 901), 78, 162, 64)
    # Часы для cutscene_4 выбираются runtime по оставшимся минутам и
    # анимируются вместе с песком; в статической основе их быть не должно.

    mouse = base.copy()
    sprite_paste(mouse, png(PV, 920), 120, 169, 64)
    sprite_paste(mouse, png(PV, 959), 18, 166, 0)

    short_time = base.copy()
    sprite_paste(short_time, png(PV, 816), 146, 166, 16)
    sprite_paste(short_time, png(PV, 953), 18, 166, 0)

    return [waiting, jaffar, magic, alone, stand, lying, mouse, short_time], colors


def pv_animation_assets() -> None:
    """Ресурсы покадрового intro.

    `b0_*` — единственный чистый фон. Он копируется accelerator-ом между
    обеими экранными страницами; спрайты рисуются банком без shadow, так что
    следующий кадр автоматически начинает с фона. Каждый actor-атлас <=16K,
    поэтому нет чтения 320x200 с HDD в анимационном цикле.
    """
    base, _ = pv_palette_canvas()
    write_scene("b", 0, base)
    # frame_tbl_cuts из SDLPoP адресует Princess как image 0..16 ресурса
    # 800 (то есть PNG res801..817). Для Jaffar PV1 адресуется ресурс 850:
    # image 0..37 = res851..888. Все 38 кадров не помещаются в одну EMM,
    # поэтому сохраняем их тремя последовательными атласами. Runtime знает
    # границы, но номер кадра остаётся исходным image-id 0..37.
    sprite_atlas(OUT / "a0.atl", sprite_frames(PV, 801, 817, 16, True))
    sprite_atlas(OUT / "j0.atl", sprite_frames(PV, 851, 863, 32, True))
    sprite_atlas(OUT / "j1.atl", sprite_frames(PV, 864, 874, 32, True))
    sprite_atlas(OUT / "j2.atl", sprite_frames(PV, 875, 888, 32, True))
    # ВТОРОЙ набор Princess (PV2, res901..930 = image 0..29 ресурса 900).
    # Он нужен сценам 8/9 (сидит и гладит мышь, встаёт, приседает) и финалу
    # (поворачивается и обнимает Кида) — поз PV1 для этого не хватает.
    # 26.7 КБ не влезают в одну EMM, поэтому два атласа; номер кадра при
    # этом остаётся исходным image-id 0..29, границу знает runtime.
    # Палитровая строка 64..79 — та же, что у PV2 в статических сценах.
    sprite_atlas(OUT / "a1.atl", sprite_frames(PV, 901, 917, 64, True))
    sprite_atlas(OUT / "a2.atl", sprite_frames(PV, 918, 930, 64, True))
    # res953/res954 — состояния часов, res960..962 — три фазы песка.
    sprite_atlas(OUT / "h0.atl", sprite_frames(PV, 952, 962, 0))
    sprite_atlas(OUT / "t0.atl", sprite_frames(PRINCE, 151, 159, 80))


def main() -> None:
    OUT.mkdir(parents=True, exist_ok=True)

    stories = story_scenes()
    write_palette(OUT / "story.pal", palette(png(TITLE, 51)),
                  {STORY_BG_INDEX: STORY_BG_COLOR,
                   STORY_FONT_INDEX: (224, 160, 0)})
    for number, scene in enumerate(stories):
        write_scene("s", number, scene)

    scenes, colors = pv_scenes()
    write_palette(OUT / "pv.pal", colors)
    for number, scene in enumerate(scenes):
        write_scene("p", number, scene)

    pv_animation_assets()

    print(f"{OUT}: {len(stories)} story + 8 PV compositions + base x 5 strips, "
          f"6 animated actor atlases, one strip/page <= {STRIP_H} rows")


if __name__ == "__main__":
    main()