#!/usr/bin/env python3
"""gen_textures.py — batch park-surface textures from MODELBEAST flux_local.

Guest token allows 4 active jobs: submit in a rolling window, poll, download
each job's own output (parent_job match — newest-asset raced once, 2026-07-16).
Token from ~/Documents/backnforth/.env (never printed). Re-runs skip textures
that already exist on disk, so the script is idempotent.

Usage: python3 tools/gen_textures.py [--force name ...]
"""
import json, os, sys, time, urllib.request

HOST = os.environ.get('MB_HOST', 'http://100.89.131.57:8777')
OUT = os.path.join(os.path.dirname(__file__), '..', 'textures')

# Every prompt asks for flat, even, shadow-free top-down macro photography —
# these get tiled with MIRRORED_REPEAT in the GLB sampler, which hides seams,
# but baked-in lighting gradients would still show. 1024px, klein-4b, 4 steps.
STYLE = ('seamless tileable texture, photographed straight down, perfectly flat, '
         'even diffuse lighting, no shadows, no vignette, full-frame surface detail, '
         'photorealistic macro material photograph')

TEXTURES = {
  'concrete_smooth': 'smooth grey skatepark concrete, steel-trowel finish, faint swirl marks, fine aggregate flecks, light wear',
  'concrete_green':  'green tinted polished skatepark concrete, smooth burnished surface, subtle colour variation, small speckle aggregate',
  'concrete_rough':  'broom finished pale concrete footpath surface, fine parallel brush lines, weathered',
  'plywood_ramp':    'worn skate ramp plywood, birch surface scuffed with black urethane wheel marks and scratches, faded grain',
  'steel_galv':      'galvanized steel sheet metal, crystalline spangle pattern, dull grey zinc coating',
  'steel_scratch':   'steel surface covered in fine bright grind scratches and wax residue, well used skateboard rail metal',
  'wood_slats':      'weathered grey-brown hardwood park bench timber, deep grain, small cracks, sun faded',
  'bark_fig':        'Moreton Bay fig tree bark, grey-brown, rough woven fibrous ridges',
  'leaves_fig':      'dense dark green glossy fig tree canopy foliage seen from below, layered leaves',
  'bark_gum':        'eucalyptus gum tree bark, smooth mottled cream grey and salmon patches, peeling strips',
  'leaves_gum':      'dense eucalyptus foliage filling the entire frame, layered slender sage-green curved gum leaves, no sky visible',
  'grass_dry':       'dry australian park grass, green-brown patchy couch lawn, some bare dirt showing',
  'tarmac':          'grey asphalt road surface, fine bitumen aggregate, small stones, light wear',
  'dirt_mulch':      'garden bed surface of brown dirt and eucalyptus bark mulch chips',
  'brick_qld':       'reddish brown queensland house bricks with pale mortar joints, running bond',
  'graffiti_a':      'colorful graffiti burner piece with bold wildstyle letters on a concrete skatepark wall, aerosol art, blues oranges purples',
  'graffiti_b':      'concrete wall covered in layered spray paint tags and simple throw-up graffiti letters, black silver white paint',
  'sky_day':         'clear bright queensland blue sky with a few small white cumulus clouds, wide even gradient',
}

def token():
    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=180).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 submit(name, desc, seed):
    p = {'operator': 'flux_local',
         'params': {'prompt': f'{desc}, {STYLE}', 'model': 'flux2-klein-4b',
                    'steps': 4, 'width': 1024, 'height': 1024, 'seed': seed}}
    j = req('/api/jobs', data=json.dumps(p).encode(),
            headers={'Content-Type': 'application/json'})
    print(f'[tx] {name}: job {j["id"]}', flush=True)
    return j['id']

def fetch(jid, name):
    items = req('/api/assets?limit=100')
    items = items if isinstance(items, list) else items.get('items', [])
    mine = [x for x in items if x.get('parent_job') == jid]
    imgs = [x for x in mine if str(x.get('filename', x.get('name', ''))).lower()
            .endswith(('.png', '.jpg', '.jpeg', '.webp'))]
    if not imgs:
        print(f'[tx] {name}: NO OUTPUT ASSET for job {jid}', flush=True)
        return False
    ext = os.path.splitext(str(imgs[0].get('filename', imgs[0].get('name', 'x.png'))))[1] or '.png'
    data = req(f'/api/assets/{imgs[0]["id"]}/file', raw=True)
    out = os.path.join(OUT, name + ext)
    open(out, 'wb').write(data)
    print(f'[tx] {name}: saved {len(data)//1024}KB -> {out}', flush=True)
    return True

def main():
    force = set(sys.argv[sys.argv.index('--force') + 1:]) if '--force' in sys.argv else set()
    todo = []
    for i, (name, desc) in enumerate(TEXTURES.items()):
        have = [f for f in os.listdir(OUT) if f.startswith(name + '.')]
        if have and name not in force:
            print(f'[tx] {name}: exists, skip', flush=True)
            continue
        todo.append((name, desc, 1000 + i))
    active = {}          # jid -> name
    fails = []
    while todo or active:
        while todo and len(active) < 4:
            name, desc, seed = todo.pop(0)
            try:
                active[submit(name, desc, seed)] = name
            except Exception as e:
                print(f'[tx] {name}: submit failed {e}', flush=True)
                fails.append(name)
        time.sleep(7)
        for jid in list(active):
            try:
                st = req(f'/api/jobs/{jid}').get('status')
            except Exception:
                continue
            if st in ('done', 'error', 'cancelled'):
                name = active.pop(jid)
                if st == 'done':
                    if not fetch(jid, name):
                        fails.append(name)
                else:
                    print(f'[tx] {name}: job {st}', flush=True)
                    fails.append(name)
    print(f'[tx] COMPLETE. fails: {fails or "none"}', flush=True)

if __name__ == '__main__':
    main()