#!/usr/bin/env python3
"""The trailer's generator: stills, clips and score through the OpenRouter key in the repo's .env.

    uv run --with requests --with pillow python tools/trailer/gen.py still <name> "<prompt>" [ref.png ...]
    uv run --with requests --with pillow python tools/trailer/gen.py clip <name> <model> <secs> <first.png> "<prompt>" [last.png]
    uv run --with requests --with pillow python tools/trailer/gen.py music <name> "<prompt>"

Everything lands in promo/trailer/ (gitignored). Each call also writes <name>.json beside its
output with the exact request, so a take can be re-rolled or compared later.
"""
import base64, json, os, sys, time

import requests

ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
OUT = os.path.join(ROOT, 'promo', 'trailer')
API = 'https://openrouter.ai/api/v1'
IMG_MODEL = os.environ.get('TRAILER_IMG', 'google/gemini-3-pro-image')


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


H = lambda: {'Authorization': f'Bearer {key()}', 'Content-Type': 'application/json'}


def u(p):
    ext = 'jpeg' if p.lower().endswith(('.jpg', '.jpeg')) else 'png'
    return f'data:image/{ext};base64,' + base64.b64encode(open(p, 'rb').read()).decode()


def note(path, req):
    req = {k: v for k, v in req.items() if k not in ('input_references', 'frame_images')}
    json.dump(req, open(os.path.splitext(path)[0] + '.json', 'w'), indent=1)


def still(name, prompt, refs=(), ar='9:16', res='2K'):
    req = {'model': IMG_MODEL, 'prompt': prompt, 'resolution': res, 'aspect_ratio': ar}
    if refs:
        req['input_references'] = [{'type': 'image_url', 'image_url': {'url': u(r)}} for r in refs]
    r = requests.post(f'{API}/images', headers=H(), json=req, timeout=1800)
    if r.status_code != 200:
        sys.exit(f'{name}: HTTP {r.status_code} {r.text[:400]}')
    d = (r.json().get('data') or [{}])[0]
    b = d.get('b64_json') or d.get('image_url', {}).get('url', '').split(',')[-1]
    if not b:
        sys.exit(f'{name}: no image returned {r.text[:400]}')
    out = os.path.join(OUT, 'stills', name + '.png')
    open(out, 'wb').write(base64.b64decode(b))
    note(out, dict(req, refs=list(refs)))
    print(out, flush=True)


def clip(name, model, secs, first, prompt, last=None, res='1080p', audio=True):
    frames = [{'type': 'image_url', 'image_url': {'url': u(first)}, 'frame_type': 'first_frame'}]
    if last:
        frames.append({'type': 'image_url', 'image_url': {'url': u(last)}, 'frame_type': 'last_frame'})
    req = {'model': model, 'prompt': prompt, 'duration': int(secs), 'aspect_ratio': '9:16',
           'resolution': res, 'generate_audio': audio, 'frame_images': frames}
    r = requests.post(f'{API}/videos', headers=H(), json=req, timeout=300)
    if r.status_code not in (200, 201, 202):
        sys.exit(f'{name}: HTTP {r.status_code} {r.text[:600]}')
    job = r.json()
    poll = job.get('polling_url') or f"{API}/videos/{job['id']}"
    t0 = time.time()
    while True:
        time.sleep(10)
        s = requests.get(poll, headers=H(), timeout=60).json()
        st = s.get('status')
        if st in ('completed', 'succeeded'):
            break
        if st in ('failed', 'error', 'cancelled'):
            sys.exit(f'{name}: {st} {json.dumps(s)[:600]}')
        if time.time() - t0 > 1500:
            sys.exit(f'{name}: timed out, job {job.get("id")}')
    urls = s.get('unsigned_urls') or s.get('urls') or [o.get('url') for o in s.get('output', []) if isinstance(o, dict)]
    out = os.path.join(OUT, 'clips', name + '.mp4')
    v = requests.get(urls[0], headers=H(), timeout=300)   # the "unsigned" URL still wants auth
    open(out, 'wb').write(v.content)
    note(out, dict(req, first=first, last=last, cost=s.get('usage') or s.get('cost')))
    print(out, len(v.content), 'bytes', f'{time.time() - t0:.0f}s', flush=True)


def music(name, prompt, model='google/lyria-3-pro-preview'):
    req = {'model': model, 'messages': [{'role': 'user', 'content': prompt}],
           'modalities': ['audio', 'text'], 'stream': True}
    r = requests.post(f'{API}/chat/completions', headers=H(), json=req, stream=True, timeout=900)
    if r.status_code != 200:
        sys.exit(f'{name}: HTTP {r.status_code} {r.text[:600]}')
    chunks, fmt, text = [], 'mp3', []
    for line in r.iter_lines():
        if not line.startswith(b'data: ') or line == b'data: [DONE]':
            continue
        ev = json.loads(line[6:])
        for ch in ev.get('choices', []):
            d = ch.get('delta', {})
            a = d.get('audio') or {}
            if a.get('data'):
                chunks.append(a['data'])
            fmt = a.get('format', fmt)
            if d.get('content'):
                text.append(d['content'])
    if not chunks:
        sys.exit(f'{name}: no audio. text: {"".join(text)[:400]}')
    out = os.path.join(OUT, 'audio', f'{name}.{fmt}')
    open(out, 'wb').write(base64.b64decode(''.join(chunks)))
    note(out, req)
    print(out, ''.join(text)[:300], flush=True)


if __name__ == '__main__':
    a = sys.argv[1:]
    if a[:1] == ['still']:
        still(a[1], a[2], a[3:])
    elif a[:1] == ['clip']:
        clip(a[1], a[2], a[3], a[4], a[5], a[6] if len(a) > 6 else None)
    elif a[:1] == ['music']:
        music(a[1], a[2], *a[3:4])
    else:
        sys.exit(__doc__)