#!/usr/bin/env python3
"""MORPQUEST art pipeline — nano banana (OpenRouter) -> Sierra EGA assets.

Usage:
  python3 tools/gen_art.py bg <slug>          one background from PROMPTS
  python3 tools/gen_art.py sprite <slug>      one magenta-keyed sprite
  python3 tools/gen_art.py all                everything not already on disk

Raw generations cached in art/raw/ (git-ignored); processed PNGs land in
sprites/ where the engine quantizes them to EGA on load (alpha<128 = clear).
"""
import base64, json, os, re, sys, urllib.request

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RAW = os.path.join(ROOT, 'art', 'raw')
SPRITES = os.path.join(ROOT, 'sprites')
os.makedirs(RAW, exist_ok=True)
os.makedirs(SPRITES, exist_ok=True)

MODEL = os.environ.get('TG_ART_MODEL', 'google/gemini-3.1-flash-image')
KEY = os.environ.get('OPENROUTER_KEY')
if not KEY:  # family convention — key lives in dealgod, never in a repo
    for line in open(os.path.expanduser('~/Documents/dealgod/lore/keys.env')):
        if line.startswith('OPENROUTER_KEY='):
            KEY = line.strip().split('=', 1)[1]

W, H = 160, 168  # engine picture resolution; GUI displays it 320x168 (2x-wide pixels)

EGA = ('strict 16-color EGA palette only (black, dark blue, dark green, dark cyan, dark red, '
       'dark magenta, brown, light gray, dark gray, bright blue, bright green, bright cyan, '
       'bright red, bright magenta, yellow, white)')

BG_STYLE = ('Background art from a 1989 Sierra On-Line EGA adventure game, King\'s Quest / '
            f'Police Quest II era. {EGA}. Flat color fills with occasional coarse dither '
            'patterns, hard aliased pixel edges, no gradients, no smooth shading, no '
            'anti-aliasing, no text, no letters, no UI, no people or animals. Wide landscape '
            'scene from a slightly raised third-person view, with an open walkable ground '
            'plane across the lower half of the frame. ')

SPRITE_STYLE = (f'A single video game sprite from a 1989 Sierra EGA adventure game. {EGA}. '
                'Flat colors, hard pixel edges, no anti-aliasing, no gradients, no text. '
                'Full figure, facing the viewer, centered, on a SOLID PURE MAGENTA '
                'background (#FF00FF) with nothing else. ')

PROMPTS = {'bg': {}, 'sprite': {}}  # filled by tools/art_manifest.py or edits below
try:
    from art_manifest import PROMPTS  # noqa: F401
except ImportError:
    pass


def gen(prompt):
    body = json.dumps({'model': MODEL, 'modalities': ['image', 'text'],
                       'messages': [{'role': 'user', 'content': prompt}]}).encode()
    req = urllib.request.Request('https://openrouter.ai/api/v1/chat/completions', data=body,
                                 headers={'Authorization': f'Bearer {KEY}',
                                          'Content-Type': 'application/json'})
    with urllib.request.urlopen(req, timeout=180) as r:
        d = json.load(r)
    msg = d['choices'][0]['message']
    imgs = msg.get('images') or []
    url = imgs[0]['image_url']['url'] if imgs else ''
    m = re.match(r'data:image/(\w+);base64,(.+)', url, re.S)
    if not m:
        raise RuntimeError(f'no image in response: {str(msg)[:200]}')
    return base64.b64decode(m.group(2))


def bg(slug, force=False):
    out = os.path.join(SPRITES, f'bg_{slug}.png')
    if os.path.exists(out) and not force:
        return out
    raw = os.path.join(RAW, f'bg_{slug}.png')
    if not os.path.exists(raw) or force:
        open(raw, 'wb').write(gen(BG_STYLE + PROMPTS['bg'][slug]))
    from PIL import Image
    img = Image.open(raw).convert('RGB')
    img = img.resize((W, H), Image.LANCZOS)  # squeeze wide gen -> 2x-wide-pixel frame
    img.save(out)
    print('bg ', slug, '->', out)
    return out


def sprite(slug, force=False):
    """Chroma-keyed sprite. PROMPTS['sprite'][slug] = (prompt, target_height_px)."""
    out = os.path.join(SPRITES, f'{slug}.png')
    if os.path.exists(out) and not force:
        return out
    prompt, target_h = PROMPTS['sprite'][slug]
    raw = os.path.join(RAW, f'sp_{slug}.png')
    if not os.path.exists(raw) or force:
        open(raw, 'wb').write(gen(SPRITE_STYLE + prompt))
    from PIL import Image
    img = Image.open(raw).convert('RGB')
    img.thumbnail((256, 256), Image.LANCZOS)  # work small: fast + smooths jpeg noise
    px = img.load()
    w, h = img.size
    # key by distance to the sampled corner colors (the model's magenta varies)
    corners = {px[0, 0], px[w-1, 0], px[0, h-1], px[w-1, h-1]}
    def bg(p):
        return any((p[0]-c[0])**2 + (p[1]-c[1])**2 + (p[2]-c[2])**2 < 3600 for c in corners)
    solid = [[not bg(px[x, y]) for x in range(w)] for y in range(h)]
    # keep only the largest connected blob (kills stray decals/debris)
    seen = [[False]*w for _ in range(h)]
    best = []
    for y0 in range(h):
        for x0 in range(w):
            if solid[y0][x0] and not seen[y0][x0]:
                blob, stack = [], [(x0, y0)]
                seen[y0][x0] = True
                while stack:
                    x, y = stack.pop()
                    blob.append((x, y))
                    for nx, ny in ((x+1,y),(x-1,y),(x,y+1),(x,y-1)):
                        if 0 <= nx < w and 0 <= ny < h and solid[ny][nx] and not seen[ny][nx]:
                            seen[ny][nx] = True
                            stack.append((nx, ny))
                if len(blob) > len(best):
                    best = blob
    alpha = Image.new('L', (w, h), 0)
    ap = alpha.load()
    for x, y in best:
        ap[x, y] = 255
    img = img.convert('RGBA')
    img.putalpha(alpha)
    img = img.crop(img.getbbox())
    # engine pixels display 2x wide; squeeze width, but keep Sierra-chunky
    ratio = target_h / img.height
    img = img.resize((max(3, round(img.width * ratio * 0.75)), target_h), Image.LANCZOS)
    img.putalpha(img.getchannel('A').point(lambda a: 255 if a >= 100 else 0))
    img.save(out)
    print('spr', slug, f'{img.width}x{img.height}', '->', out)
    return out


if __name__ == '__main__':
    mode = sys.argv[1] if len(sys.argv) > 1 else 'all'
    if mode == 'bg':
        bg(sys.argv[2], force='-f' in sys.argv)
    elif mode == 'sprite':
        sprite(sys.argv[2], force='-f' in sys.argv)
    else:
        for s in PROMPTS['bg']:
            bg(s)
        for s in PROMPTS['sprite']:
            sprite(s)