#!/usr/bin/env python3
"""restyle_frames.py — generate frames@<style>/ variant dirs for a packed
character by restyling its existing base frames. Two engines:

  pillow-xray   deterministic local pass (no AI): dark silhouette fill +
                glowing edge lines. Instant, zero flicker — the "digital
                guts" tier. Needs Pillow only.
  farm          any MODELBEAST image->image operator (mflux_image_edit,
                comfyui_sd, ...). Frame is flattened onto neutral gray,
                submitted, and the result is re-masked with the original
                alpha so the sprite silhouette is untouched. Fixed seed
                per style keeps frame-to-frame drift down (some shimmer
                is period-correct for digitized fighters). SLOW operators
                (mflux_image_edit ≈ 4 min/frame on m3ultra) are hero-tier
                only — check timing on one frame before a big batch.

Usage:
  python3 pipeline/restyle_frames.py characters/vesper --style xray --engine pillow-xray
  python3 pipeline/restyle_frames.py characters/vesper --style sketch --engine farm \
      --operator mflux_image_edit --prompt "Redraw as a pencil sketch ..." \
      --op-params '{"steps": 20}' --moves idle,hit --limit 10

Resume-safe: frames whose output already exists are skipped. Variant dirs
must keep the SAME frame count as base (engine falls back per-frame past
the variant's end, which pops mid-move) — this script always writes 1:1.
Heartbeat: --hb writes progress lines to ~/.jobs/foitin-restyle.hb.
"""
import argparse, json, mimetypes, os, sys, time, urllib.request, uuid

HOST = os.environ.get('MB_HOST', 'http://100.89.131.57:8777')
HB = os.path.expanduser('~/.jobs/foitin-restyle.hb')


def hb(msg, enabled):
    if not enabled:
        return
    os.makedirs(os.path.dirname(HB), exist_ok=True)
    with open(HB, 'a') as f:
        f.write(time.strftime('%H:%M:%S ') + msg + '\n')


# ---------------------------------------------------------------- farm client

def token():
    t = os.environ.get('MB_TOKEN')
    if t:
        return t
    for line in open(os.path.expanduser('~/Documents/backnforth/.env')):
        if line.startswith('MB_TOKEN='):
            return line.split('=', 1)[1].strip()
    sys.exit('no MB_TOKEN')


def req(path, data=None, headers=None, raw=False):
    h = {'Authorization': f'Bearer {token()}'}
    h.update(headers or {})
    r = urllib.request.Request(HOST + path, data=data, headers=h)
    body = urllib.request.urlopen(r, timeout=300).read()
    if raw:
        return body
    s = body.decode('utf-8', 'replace')
    return json.loads(''.join(c if c >= ' ' or c in '\t' else ' ' for c in s))


def upload(path):
    boundary = uuid.uuid4().hex
    ctype = mimetypes.guess_type(path)[0] or 'application/octet-stream'
    body = (f'--{boundary}\r\nContent-Disposition: form-data; name="file"; '
            f'filename="{os.path.basename(path)}"\r\nContent-Type: {ctype}\r\n\r\n').encode() \
        + open(path, 'rb').read() + f'\r\n--{boundary}--\r\n'.encode()
    a = req('/api/assets', data=body, headers={'Content-Type': f'multipart/form-data; boundary={boundary}'})
    return a.get('id') or (a.get('items') or [a])[0].get('id')


def job_output_asset(jid):
    a = req('/api/assets?limit=100')
    items = a if isinstance(a, list) else a.get('items', [])
    imgs = [x for x in items if str(x.get('filename', x.get('name', ''))).lower()
            .endswith(('.png', '.jpg', '.webp'))]
    mine = [x for x in imgs if x.get('parent_job') == jid]
    return mine[0]['id'] if mine else None


def farm_edit(Image, src_png, out_png, operator, params, tmp_dir, control='flat'):
    im = Image.open(src_png).convert('RGBA')
    alpha = im.split()[3]
    tmp_in = os.path.join(tmp_dir, 'mb_in.png')
    if control == 'edge':
        # canny-style conditioning image (thin white lines on black) for
        # ControlNet graphs with no preprocessor node (comfyui_sd)
        from PIL import ImageChops, ImageFilter, ImageOps
        edges = ImageOps.autocontrast(im.convert('L').filter(ImageFilter.FIND_EDGES))
        ImageChops.multiply(edges, alpha).convert('RGB').save(tmp_in)
    else:
        flat = Image.new('RGB', im.size, (128, 128, 128))
        flat.paste(im, (0, 0), alpha)
        flat.save(tmp_in)
    jid = req('/api/jobs', data=json.dumps(
        {'operator': operator, 'asset_id': upload(tmp_in), 'params': params}).encode(),
        headers={'Content-Type': 'application/json'})['id']
    while True:
        time.sleep(5)
        st = req(f'/api/jobs/{jid}').get('status')
        if st in ('done', 'error', 'cancelled'):
            break
    if st != 'done':
        return False
    gid = job_output_asset(jid)
    if not gid:
        return False
    tmp_out = os.path.join(tmp_dir, 'mb_out.png')
    open(tmp_out, 'wb').write(req(f'/api/assets/{gid}/file', raw=True))
    edited = Image.open(tmp_out).convert('RGB')
    if edited.size != im.size:
        edited = edited.resize(im.size, Image.LANCZOS)
    edited = edited.convert('RGBA')
    edited.putalpha(alpha)
    edited.save(out_png)
    return True


# ------------------------------------------------------------- pillow engine

def pillow_xray(Image, src_png, out_png, fill=(10, 22, 34), line=(87, 230, 255)):
    from PIL import ImageChops, ImageFilter, ImageOps
    im = Image.open(src_png).convert('RGBA')
    alpha = im.split()[3]
    gray = im.convert('L')
    edges = ImageOps.autocontrast(gray.filter(ImageFilter.FIND_EDGES))
    edges = edges.filter(ImageFilter.MaxFilter(3))
    edges = ImageChops.multiply(edges, alpha)   # clip glow to the silhouette
    out = Image.new('RGBA', im.size, (0, 0, 0, 0))
    out.paste(Image.new('RGBA', im.size, fill + (255,)), (0, 0), alpha)
    out.paste(Image.new('RGBA', im.size, line + (255,)), (0, 0), edges)
    out.putalpha(alpha)
    out.save(out_png)
    return True


# ----------------------------------------------------------------------- main

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('char_dir', help='packed character folder (characters/<id>)')
    ap.add_argument('--style', required=True, help='variant name -> frames@<style>/')
    ap.add_argument('--engine', default='pillow-xray', choices=['pillow-xray', 'farm'])
    ap.add_argument('--operator', default='mflux_image_edit')
    ap.add_argument('--prompt', default='')
    ap.add_argument('--op-params', default='{}', help='extra operator params as JSON')
    ap.add_argument('--control', default='flat', choices=['flat', 'edge'],
                    help='what to upload as the job asset: flattened frame, or an edge map (controlnet)')
    ap.add_argument('--moves', default='', help='comma list; default = all moves')
    ap.add_argument('--limit', type=int, default=0, help='max frames this run (0 = all)')
    ap.add_argument('--hb', action='store_true', help='heartbeat to ~/.jobs/foitin-restyle.hb')
    a = ap.parse_args()

    from PIL import Image
    params = json.loads(a.op_params)
    if a.prompt:
        params['prompt'] = a.prompt
    params.setdefault('seed', 42)

    moves_root = os.path.join(a.char_dir, 'moves')
    moves = sorted(a.moves.split(',') if a.moves else os.listdir(moves_root))
    tmp_dir = f'/tmp/restyle_{os.getpid()}'
    os.makedirs(tmp_dir, exist_ok=True)

    done = failed = 0
    for mv in moves:
        frames_dir = os.path.join(moves_root, mv, 'frames')
        if not os.path.isdir(frames_dir):
            continue
        out_dir = os.path.join(moves_root, mv, 'frames@' + a.style)
        os.makedirs(out_dir, exist_ok=True)
        pngs = sorted(f for f in os.listdir(frames_dir) if f.endswith('.png'))
        for i, f in enumerate(pngs):
            out_png = os.path.join(out_dir, f)
            if os.path.exists(out_png):
                continue
            if a.limit and done >= a.limit:
                hb(f'limit {a.limit} reached', a.hb)
                print(f'[restyle] limit reached ({done} frames)')
                return
            src = os.path.join(frames_dir, f)
            if a.engine == 'pillow-xray':
                ok = pillow_xray(Image, src, out_png)
            else:
                ok = farm_edit(Image, src, out_png, a.operator, dict(params), tmp_dir, a.control)
            done += ok
            failed += (not ok)
            if not ok:
                print(f'[restyle] FAILED {mv}/{f}', flush=True)
            if done % 20 == 0 or a.engine == 'farm':
                hb(f'{mv} {i + 1}/{len(pngs)} (total {done} done {failed} failed)', a.hb)
                print(f'[restyle] {mv} {i + 1}/{len(pngs)} (total {done})', flush=True)
    hb(f'ALL DONE ({done} frames, {failed} failed)', a.hb)
    print(f'[restyle] ALL DONE ({done} frames, {failed} failed)')


main()