#!/usr/bin/env python3
"""AI-шаг пробы B1: доменные интерактивы через PixelGPT-24×24 (CPU).

Модель: unstonio/pixelgpt-24x24 (68M декодер, ~22 с/спрайт). Репо модели —
вне репо игры: ~/.cache/rpg-ai/pixelgpt-24x24 (env RPG_PXGPT_REPO). Веса —
models/pixelar_fp16.pt (LFS) + models/minilm/; лицензии на модель нет —
результаты только PoC/review, в игру не идут.

Палитра обусловливания — 5 ключей из apps/game/tools/pixelart/palette.mjs
(первый — самый тёмный, токен 0 = прозрачный). Токены 1–4 маппятся на
палитру по индексу (как во фронтенде модели: fillStyle = palette[role]).

Запуск: $RPG_SD_PYTHON gen_pxgpt.py <prompts-файл> <outdir>
  строки промптов: seed|имя|промпт|5 ключей палитры через запятую
  (формат apps/game/tools/aiart/prompts-b1.txt); строки с # — комментарии.
"""
import os
import re
import sys
import time
from pathlib import Path

import numpy as np
import torch
from PIL import Image

REPO = Path(os.environ.get('RPG_PXGPT_REPO', f'{Path.home()}/.cache/rpg-ai/pixelgpt-24x24'))
GAME_ROOT = Path(__file__).resolve().parents[4]
CFG = 3.0
TOP_P = 0.95
TEMP = 1.0


def load_game_palette(path: Path) -> dict[str, tuple[int, int, int]]:
    """Ключи палитры игры → RGB (парсинг palette.mjs)."""
    out = {}
    for m in re.finditer(r"(?:^|\s)([A-Z]\w*):\s*'(#[0-9a-fA-F]{6})'", path.read_text()):
        h = m.group(2)
        out[m.group(1)] = tuple(int(h[i:i + 2], 16) for i in (1, 3, 5))
    return out


def render(tokens: list[int], palette: list[tuple[int, int, int]], out: Path) -> None:
    """24×24 RGBA: токен 0 — прозрачный, 1–5 — палитра со сдвигом (t-1).

    Сдвиг восстановлен сверкой с review-b1: токен 1 = первый ключ палитры
    (самый тёмный), токен 5 — последний (в пробы B1 не попадал).
    """
    im = Image.new('RGBA', (24, 24), (0, 0, 0, 0))
    px = im.load()
    for i, t in enumerate(tokens):
        if t:
            px[i % 24, i // 24] = (*palette[t - 1], 255)
    out.parent.mkdir(parents=True, exist_ok=True)
    im.save(out)


def main() -> None:
    prompts_file, out_dir = Path(sys.argv[1]), Path(sys.argv[2])
    keys = load_game_palette(GAME_ROOT / 'apps/game/tools/pixelart/palette.mjs')

    sys.path.insert(0, str(REPO))
    torch.set_num_threads(4)
    from model_runtime import SEQ_LEN, SOS_TOKEN, load_model, sample_next, seed_everything, sequence_to_tokens
    from sentence_transformers import SentenceTransformer

    device = torch.device('cpu')
    t0 = time.time()
    model, config, flat = load_model(str(REPO / 'models/pixelar_fp16.pt'), device)
    encoder = SentenceTransformer(str(REPO / 'models/minilm'), device='cpu', local_files_only=True)
    print(f'[load] {time.time() - t0:.0f}s', flush=True)

    for line in prompts_file.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith('#'):
            continue
        seed, name, prompt, palette_keys = line.split('|')
        palette = [keys[k.strip()] for k in palette_keys.split(',')]
        seed = int(seed)
        seed_everything(seed)
        dtype = next(model.parameters()).dtype
        embedding = torch.from_numpy(encoder.encode(
            [prompt], normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False,
        )).to(device=device, dtype=dtype)
        pal = torch.from_numpy(np.asarray(palette, dtype=np.uint8)[None]).to(device=device, dtype=dtype) / 255.0
        captions = torch.cat([embedding, model.null_caption.unsqueeze(0).to(embedding.dtype)], dim=0)
        model_palettes = pal.repeat(2, 1, 1)
        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=False):
            condition = model.cond_vector(captions, model_palettes)
            caches = model.allocate_caches(2, device, condition.dtype)
            current = torch.full((2, 1), SOS_TOKEN, dtype=torch.long, device=device)
            sequence = torch.zeros(1, SEQ_LEN, dtype=torch.long, device=device)
            for position in range(SEQ_LEN):
                logits = model(current, model_palettes, condition, caches=caches, start_pos=position)[:, -1].float()
                conditional, unconditional = logits[:1], logits[1:]
                token = sample_next(unconditional + CFG * (conditional - unconditional), TEMP, TOP_P)
                sequence[:, position] = token
                current = torch.cat([token, token], dim=0).unsqueeze(1)
        tokens = sequence_to_tokens(sequence, flat).clamp(0, 4)[0].cpu().to(torch.uint8).reshape(-1).tolist()
        render(tokens, palette, out_dir / f'{name}.png')
        print(f'[ok] {name} (seed {seed}, {time.time() - t0:.0f}s)', flush=True)


if __name__ == '__main__':
    main()