# _frontroll.py — the NETPLAY OPPONENT: the player's own kit, drawn from the FRONT.
#
#   uv run --with pillow --with numpy --with requests python assets/_frontroll.py m --ref
#   uv run --with pillow --with numpy --with requests python assets/_frontroll.py m
#   uv run --with pillow --with numpy python assets/_frontroll.py m --cut          # review
#   uv run --with pillow --with numpy python assets/_frontroll.py m --cut --install
#
# WHY THIS EXISTS. In a netplay match the body across the net is your FRIEND, and it was drawn on
# the roster opponent's sheets: red polo, dark grey shorts, red headband. His pick then dyed the
# only thing on that art that carries a hue — the SHIRT — while your own pick dyes your SHORTS,
# because the player's kit is white-over-blue. Two people in one match, wearing their colour on
# different garments. Worse on the women's sheets: the opponent's red sits at the same hue as her
# skin and her hair, so `TINT.oKit` bled green over her arms, her legs and her ponytail.
#
# Neither is fixable in the tint pass — the roster opponent's art simply has no hue-isolated kit —
# so the friend gets his own sheets, `n_*`, in the PLAYER'S kit: white shirt, blue shorts, and all
# whites for the women's tour, which is exactly what `TINT.kit`/`wash` are measured against.
#
# THE METHOD IS `_famroll.py`'s, and it is the reason this is two calls and not ten. IMAGE 1 is the
# POSES (the installed o_* frames, composed from disk at one scale), IMAGE 2 is the CHARACTER (the
# player, from behind) — "redraw image 1 with one change: the character". Ten separate _regen
# batches would be ten chances for the man to drift; the whole family in one roll cannot drift.
import base64, os, sys
import numpy as np
import requests
from PIL import Image, ImageDraw

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from _famref import bed_magenta                       # the racket bed has to key out — see _zsheet
from _zsheet import cut

HERE = os.path.dirname(os.path.abspath(__file__))
SCRATCH = os.path.join(HERE, '_refcheck')
MODEL = os.environ.get('SPLAT_MODEL', 'google/gemini-3-pro-image')
MAGENTA = (255, 0, 255)
SIZE = (4304, 5380)                                   # 4:5, the enum value nearest this sheet's own aspect
MARGIN, GAP = 0.025, 60

# 4 / 4 / 2, which lands within a hair of 4:5. The pairs that MUST come back the same height sit
# next to each other so the levelling below reads them off one row.
ROWS = [['idle1', 'idle2', 'bounce', 'bounce2'],
        ['swing1', 'swing2', 'toss', 'smash'],
        ['run1', 'run2']]
# ...and these are the ones the game flips between at speed: a 2 px height difference in a pair is
# a body that hops once per bounce or once per stride (the lesson _zsheet's walk levelling records).
# The serve toss and the smash are NOT levelled — they are meant to be taller.
LEVEL = [['idle1', 'idle2'], ['bounce', 'bounce2'], ['run1', 'run2']]

WHO = {
 '': ('a male tennis player in a PURE WHITE polo shirt, ROYAL BLUE shorts, a WHITE cloth headband, '
      'white socks and white tennis shoes, short blond hair, lean and athletic. His shirt is WHITE, '
      'never red. His shorts are BRIGHT ROYAL BLUE, never grey, never black, never navy. His '
      'headband and his socks are WHITE, never red'),
 '_f': ('a female tennis player dressed head to hem in TENNIS WHITES: a PURE WHITE sleeveless-cut '
        'tennis top and a PURE WHITE pleated tennis skirt, white socks and white tennis shoes, dark '
        'hair tied back in a ponytail. Her top is WHITE, never red. Her skirt is WHITE, never dark '
        'grey and never black. Every garment she wears is white'),
}


def open_bed(im):
    """Transparent pixels the border cannot reach — i.e. the racket's string bed, already keyed out.

    `bed_magenta` repaints the largest PALE blob, which is right for the z_r* master (its beds are
    painted) and wrong for most of these: with the bed already open the largest pale blob is a SHOE,
    and the reference went out with a magenta boot on it. Ask first.
    """
    a = np.array(im.convert('RGBA'))[..., 3] > 40
    m = Image.new('L', (im.width + 2, im.height + 2), 0)
    m.paste(Image.fromarray((a * 255).astype('uint8')), (1, 1))
    ImageDraw.floodfill(m, (0, 0), 128)
    return int((np.array(m)[1:-1, 1:-1] == 0).sum())


def compose(sfx):
    """IMAGE 1: the ten tennis poses, composed from the installed sprites at ONE scale. No API call."""
    frames = {}
    for r in ROWS:
        for n in r:
            im = Image.open(os.path.join(HERE, f'o_{n}{sfx}.png')).convert('RGBA')
            frames[n] = im if open_bed(im) > 100 else bed_magenta(im)
    row_w = [sum(frames[n].width for n in r) + GAP * (len(r) - 1) for r in ROWS]
    row_h = [max(frames[n].height for n in r) for r in ROWS]
    cw, ch = max(row_w), sum(row_h) + GAP * (len(ROWS) - 1)
    k = min(SIZE[0] * (1 - 2 * MARGIN) / cw, SIZE[1] * (1 - 2 * MARGIN) / ch)
    sheet = Image.new('RGB', SIZE, MAGENTA)
    y = round((SIZE[1] - ch * k) / 2)
    for row, rh in zip(ROWS, row_h):
        scaled = [(n, round(frames[n].width * k), round(frames[n].height * k)) for n in row]
        x = round((SIZE[0] - (sum(w for _, w, _ in scaled) + round(GAP * k) * (len(row) - 1))) / 2)
        ground = y + round(rh * k)                    # one ground line per band, bottom-aligned
        for n, w, h in scaled:
            im = frames[n].resize((w, h), Image.NEAREST)
            sheet.paste(im, (x, ground - h), im)
            x += w + round(GAP * k)
        y = ground + round(GAP * k)
    p = os.path.join(SCRATCH, f'front_ref{sfx or "_m"}.png')
    os.makedirs(SCRATCH, exist_ok=True)
    sheet.save(p)
    print(f'{p}  {sheet.size}  scale {k:.3f}  10 figures, {len(ROWS)} bands')
    return p


def char_ref(sfx):
    """IMAGE 2: the player, at 4x NEAREST on magenta. Raw sprite pixels — the famroll rule."""
    im = Image.open(os.path.join(HERE, f'p_idle1{sfx}.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), MAGENTA)
    out.paste(big, (pad, pad), big)
    p = os.path.join(SCRATCH, f'front_char{sfx or "_m"}.png')
    out.save(p)
    return p


def prompt(sfx):
    her = sfx == '_f'
    return (
     'Take IMAGE 1 and redraw it with ONE change: the person wearing the poses.\n\n'
     'IMAGE 1 is a sprite sheet of a tennis player on a flat magenta background: TEN figures in '
     'THREE rows — a waiting pair, a ball-bouncing pair, a forehand backswing and follow-through, a '
     'serve toss, a smash, and a running pair.\n\n'
     f'IMAGE 2 is THE PERSON to draw instead: {WHO[sfx]}. Copy that build and every one of those kit '
     'colours exactly. IMAGE 2 shows this person FROM BEHIND; every figure you draw is seen FROM '
     'THE FRONT, facing the viewer, exactly as image 1\'s figures do. You are inventing the face.\n\n'
     'Keep IMAGE 1 EXACTLY as it is in every other respect: the same 3 rows, the same number of '
     'figures per row, the same order, the same poses, the same facing, the same sizes, the same '
     'ground lines. Only the person changes. Do NOT carry over image 1\'s RED shirt, RED headband, '
     'RED socks or DARK GREY shorts — no red and no dark grey garment appears anywhere in your '
     'output.\n\n'
     # the bounce pair is one drawing with the hand moved; the game flips between them at 3 Hz and
     # any second difference reads as a flicker. Image 1 already has this property — say so.
     'The two BALL-BOUNCING figures at the end of row 1 are the SAME drawing twice and differ ONLY '
     'in the height of the free hand. Keep them identical in every other pixel — same stance, same '
     'legs, same head, same racket, same size — exactly as image 1 has them.\n\n'
     # the game draws the ball itself, and a drawn one lands in the sprite as a blob stuck to his
     # hand. Image 1 has none; the first roll added one to both bouncing frames anyway.
     'THERE IS NO TENNIS BALL ANYWHERE IN THE PICTURE. Do not draw a ball in any figure, not in the '
     'bouncing frames, not in the serve, not in the air, not on the ground. No motion lines.\n\n'
     'EVERY FIGURE STILL HOLDS THE TENNIS RACKET image 1 gives it, in the same hand, in the same '
     'frames — the full racket with its head, its strings, its throat and its grip.\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 dark 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'
     # rule 15: the keyer deletes every pixel where red AND blue both beat green
     'NOTHING in any figure is pink, magenta, violet or purple.\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.'
     + (' Her skin is warm and human — it is never white and never grey; only her CLOTHES are white.'
        if her else ''))


def key_env():
    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 u(p):
    return 'data:image/png;base64,' + base64.b64encode(open(p, 'rb').read()).decode()


def roll(sfx):
    ref, ch = compose(sfx), char_ref(sfx)
    r = requests.post('https://openrouter.ai/api/v1/images',
        headers={'Authorization': f'Bearer {key_env()}', 'Content-Type': 'application/json'},
        json={'model': MODEL, 'prompt': prompt(sfx), 'resolution': '4K', 'aspect_ratio': '4:5',
              '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:
        sys.exit(f'HTTP {r.status_code} {r.text[:300]}')
    d = r.json().get('data') or []
    if not d:
        sys.exit('no image returned')
    b = d[0].get('b64_json') or d[0].get('image_url', {}).get('url', '').split(',')[-1]
    out = sheet_path(sfx)
    open(out, 'wb').write(base64.b64decode(b))
    print(f'{out}  {Image.open(out).size}', flush=True)


def sheet_path(sfx):
    return os.path.join(SCRATCH, f'front{sfx or "_m"}.png')


if __name__ == '__main__':
    sfx = {'m': '', 'w': '_f'}.get((sys.argv[1:2] or [''])[0])
    if sfx is None: sys.exit('usage: _frontroll.py m|w [--ref] [--cut [--install]]')
    if '--cut' in sys.argv:
        want = [[f'n_{n}{sfx}' for n in r] for r in ROWS]
        level = [[f'n_{n}{sfx}' for n in g] for g in LEVEL]
        out = HERE if '--install' in sys.argv else os.path.join(SCRATCH, f'front_cut{sfx or "_m"}')
        cut(sheet_path(sfx), None, out, want=want, level=level)
    elif '--ref' in sys.argv:
        compose(sfx); char_ref(sfx)
    else:
        roll(sfx)