# _famroll.py — roll ONE opponent family: 31 sprites, one call, one character.
#
#   uv run --with pillow --with requests python assets/_famroll.py br      # one family
#   uv run --with pillow --with requests python assets/_famroll.py sh bb   # several, in sequence
#
# Then cut it back down, review first and install second:
#   uv run --with pillow --with numpy python assets/_zsheet.py assets/_refcheck/fam_br.png br --out assets/_refcheck/sprites_br
#   uv run --with pillow --with numpy python assets/_zsheet.py assets/_refcheck/fam_br.png br --install
#
# WHY THIS IS IN THE REPO. Every family in ZFAMS except the club pro was rolled by a script that lived
# in a session scratchpad and is gone, off a reference in a gitignored directory. The pipeline PROJECT.md
# calls the one reliable way to make this game's art was, in practice, unreproducible. This is that
# script, tracked, reading the tracked reference that `_famref.py` composes.
#
# IMAGE 1 IS THE POSES, IMAGE 2 IS THE CHARACTER, and naming each image's job in the prompt is
# load-bearing: without it the majority wins and the roll comes back as whatever most of the pixels
# look like. See PROJECT.md, "What the model will and will not do".
import base64, os, sys
import requests
from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
REF = os.path.join(HERE, '_pose_ref_family.png')
SCRATCH = os.path.join(HERE, '_refcheck')
MODEL = os.environ.get('SPLAT_MODEL', 'google/gemini-3-pro-image')

# key -> (the frame the character is read off, the words that describe it). The frame is the source of
# truth for build and colour; the words only have to stop the model averaging it into a generic zombie.
FAMS = {
 'br': ('z_g1', 'an ENORMOUS FAT ZOMBIE BRUTE: a huge sagging belly over the belt, a filthy pale '
                'grey-green work shirt straining over it, heavy olive-brown work trousers, thick '
                'tree-trunk legs, scuffed brown boots, and a small bald head sunk low into massive '
                'rounded shoulders. Far bulkier, heavier and wider than the figure in image 1'),
 'sh': ('z_a1', 'a lean SHAMBLING ZOMBIE in a filthy dark red tennis shirt and grey shorts, gaunt and stooped'),
 'bb': ('z_e1', 'a small young ZOMBIE BALL BOY in a pale tournament polo and shorts, slight and short'),
 'st': ('z_d1', 'a ZOMBIE STEWARD in a dark uniform jacket with a high collar, stiff and official'),
 'ru': ('z_f1', 'a wiry FAST ZOMBIE RUNNER in a sleeveless tennis top, lean and forward-leaning'),
 'gk': ('z_boss1', 'a heavy ZOMBIE GROUNDSKEEPER in stained work overalls and a shapeless hat'),
 'lj': ('z_judge1', 'a ZOMBIE LINE JUDGE in a dark blazer and pale trousers, upright and severe'),
}


def key():
    if os.environ.get('OPENROUTER_API_KEY'): return os.environ['OPENROUTER_API_KEY']
    for line in open(os.path.join(HERE, '..', '.env')):
        if line.startswith('OPENROUTER_API_KEY='):
            return line.split('=', 1)[1].strip().strip('"').strip("'")
    sys.exit('no OPENROUTER_API_KEY')


def char_ref(k, spr):
    """IMAGE 2: the character, at 4x NEAREST on magenta. Raw sprite pixels, no round trip — the same
    rule the pose sheet now follows, and for the same reason."""
    im = Image.open(os.path.join(HERE, spr + '.png')).convert('RGBA')
    big = im.resize((im.width * 4, im.height * 4), Image.NEAREST)
    pad = round(big.height * 0.08)
    out = Image.new('RGB', (big.width + 2 * pad, big.height + 2 * pad), (255, 0, 255))
    out.paste(big, (pad, pad), big)
    p = os.path.join(SCRATCH, f'char_{k}.png')
    os.makedirs(SCRATCH, exist_ok=True)
    out.save(p)
    return p


def prompt(who):
    return (
     'Take IMAGE 1 and redraw it with ONE change: the character.\n\n'
     'IMAGE 1 is a sprite sheet of a zombie tennis player: 31 figures in 9 rows on a flat magenta '
     'background, first undamaged, then with one arm torn off, then with the head gone.\n\n'
     f'IMAGE 2 is THE CHARACTER to draw instead: {who}. It is a walking corpse with dead '
     'grey-green skin. Copy its build, its garments and every one of its colours exactly.\n\n'
     'Keep IMAGE 1 EXACTLY as it is in every other respect: the same 9 rows, the same number of '
     'figures per row, the same order, the same poses, the same facing, the same sizes, the same '
     'damage states in the same places. A figure with no head keeps no head; a figure missing an arm '
     'keeps that arm missing with the same stump and blood. Only the body wearing the poses changes. '
     'Do not carry over the previous character\'s shirt, shorts, shoes, hair or headband.\n\n'
     # EVERY UNDAMAGED FIGURE HAS A HEAD, said out loud. The clause above ("a figure with no head
     # keeps no head") is there for the wound rows and the old normalisation pass took it as licence
     # to remove one from the play row — which is how six families came back with a decapitated
     # shuffle frame. The reference no longer shows that; the prompt no longer implies it either.
     'EVERY FIGURE IN THE FIRST THREE ROWS HAS A HEAD ON ITS SHOULDERS, and so does every figure in '
     'rows 4, 5 and 6. Only the last three rows are headless. Never remove a head that image 1 shows.\n\n'
     'MATCH THE DRAWING STYLE OF IMAGE 1 EXACTLY. Image 1 is hard-edged pixel art and your output '
     'must be indistinguishable from it in technique: the same square pixel grid at the same pixel '
     'size, the same hard black outline around every figure, the same flat blocks of colour, the '
     'same stepped aliased edges, the same hand-placed dithering. No soft brushwork, no airbrushing, '
     'no blur, no gradients, no anti-aliasing.\n\n'
     'The background is a solid, flat, uniform chroma key. Use EXACTLY hex color #FF00FF, one single '
     'flat tone across the whole canvas behind every figure, with no panels or boxes of any other '
     'colour. The inside of every racket head is that same EXACT #FF00FF. No text, no shadows, '
     'no ground lines.')


def u(p):
    return 'data:image/png;base64,' + base64.b64encode(open(p, 'rb').read()).decode()


def roll(k):
    spr, who = FAMS[k]
    ch = char_ref(k, spr)
    r = requests.post('https://openrouter.ai/api/v1/images',
        headers={'Authorization': f'Bearer {key()}', 'Content-Type': 'application/json'},
        # 4K and 9:16: the endpoint DEFAULTS to 1K and snaps the aspect to an enum, so both are
        # spelled out. The reference is composed at 9:16 for exactly this reason.
        json={'model': MODEL, 'prompt': prompt(who), 'resolution': '4K', 'aspect_ratio': '9:16',
              'input_references': [{'type': 'image_url', 'image_url': {'url': u(REF)}},
                                   {'type': 'image_url', 'image_url': {'url': u(ch)}}]},
        timeout=1800)
    if r.status_code != 200:
        print(f'{k}: HTTP {r.status_code} {r.text[:200]}'); return
    d = r.json().get('data') or []
    if not d:
        print(f'{k}: no image returned'); return
    b = d[0].get('b64_json') or d[0].get('image_url', {}).get('url', '').split(',')[-1]
    out = os.path.join(SCRATCH, f'fam_{k}.png')
    open(out, 'wb').write(base64.b64decode(b))
    print(f'{k}: {out} {Image.open(out).size}  (character from {spr})', flush=True)


if __name__ == '__main__':
    keys = [a for a in sys.argv[1:] if not a.startswith('-')]
    if not keys: sys.exit(f'usage: _famroll.py <key>...   keys: {" ".join(FAMS)}')
    for k in keys:
        if k not in FAMS: sys.exit(f'unknown family {k}')
        roll(k)