#!/usr/bin/env python3
"""API429 image queue example. Dry run by default. Python 3.10+, macOS/Linux."""
import argparse
import base64
import concurrent.futures
import fcntl
import hashlib
import json
import os
from pathlib import Path
import urllib.error
import urllib.request

ENDPOINT = 'https://gateway.api429.com/v1/images/generations'

def write_atomic(path, data):
    temp = path.with_suffix(path.suffix + '.tmp')
    with open(temp, 'wb') as f:
        f.write(data)
        f.flush()
        os.fsync(f.fileno())
    os.replace(temp, path)

def save_state(folder, state):
    write_atomic(folder / 'state.json', json.dumps(state).encode())

def image_extension(data):
    if data.startswith(b'\x89PNG\r\n\x1a\n'):
        return 'png'
    if data.startswith(b'\xff\xd8\xff'):
        return 'jpg'
    if data[:4] == b'RIFF' and data[8:12] == b'WEBP':
        return 'webp'
    raise ValueError('Unsupported image signature; inspect saved response')

def generate(folder, payload, api_key, opener=urllib.request.urlopen):
    # A folder represents one immutable task, NOT server-side idempotency.
    state_path = folder / 'state.json'
    if state_path.exists():
        prior = json.loads(state_path.read_text())
        return {'id': folder.name, 'state': prior['state'], 'skipped': True}
    state = {'id': folder.name, 'state': 'in_flight'}
    save_state(folder, state)  # Persist BEFORE sending a paid request.
    req = urllib.request.Request(ENDPOINT, data=json.dumps(payload).encode(),
        headers={'Authorization': 'Bearer ' + api_key, 'Content-Type': 'application/json'})
    try:
        with opener(req, timeout=180) as response:
            status = response.status
            raw = response.read()
        write_atomic(folder / 'response.json', raw)
        state['http_status'] = status
        if status == 202:
            state['state'] = 'accepted'  # Retrieve this job; do not submit again.
        elif status != 200:
            state['state'] = 'needs_review'
        else:
            body = json.loads(raw)
            items = body.get('data', [])
            if len(items) != 1 or not items[0].get('b64_json'):
                raise ValueError('Expected one base64 image; inspect saved response')
            data = base64.b64decode(items[0]['b64_json'], validate=True)
            ext = image_extension(data)
            output = folder / ('image.' + ext)
            write_atomic(output, data)
            state.update(state='saved', file=output.name,
                         sha256=hashlib.sha256(data).hexdigest())
    except urllib.error.HTTPError as error:
        state.update(state='needs_review', http_status=error.code,
                     retry_after=error.headers.get('Retry-After'))
        write_atomic(folder / 'response-error.txt', error.read())
    except Exception as error:
        # Timeout, disconnect, malformed response or disk failure: outcome unknown.
        # Never print request headers / credentials or auto-resubmit.
        state.update(state='unknown', error_type=type(error).__name__)
    save_state(folder, state)
    return state

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('input', help='JSONL rows: {"id":"unique-variant-id","prompt":"..."}')
    parser.add_argument('--out', default='image-run')
    parser.add_argument('--limit', type=int, default=3)
    parser.add_argument('--workers', type=int, choices=range(1, 5), default=2)
    parser.add_argument('--execute', action='store_true', help='Send paid API requests')
    args = parser.parse_args()
    if not 1 <= args.limit <= 10000:
        parser.error('--limit must be 1..10000')
    rows = [json.loads(line) for line in Path(args.input).read_text().splitlines() if line.strip()]
    ids = [r.get('id') for r in rows]
    if any(not isinstance(i, str) or not i for i in ids) or len(ids) != len(set(ids)):
        parser.error('IDs must be unique nonempty strings')
    if any(not isinstance(r.get('prompt'), str) or not r['prompt'].strip() for r in rows):
        parser.error('Each row needs a nonempty prompt')
    tasks = []
    for row in rows[:args.limit]:
        payload = dict(model='gemini-3.1-flash-image', prompt=row['prompt'], n=1,
                       size='1024x1024', response_format='b64_json')
        identity = json.dumps([row['id'], payload], sort_keys=True, ensure_ascii=False).encode()
        tasks.append((hashlib.sha256(identity).hexdigest(), payload))
    print(json.dumps({'tasks': len(tasks), 'workers': args.workers,
                      'mode': 'paid' if args.execute else 'dry_run'}))
    if not args.execute:
        return
    api_key = os.environ.get('API429_API_KEY')
    if not api_key:
        parser.error('Set API429_API_KEY in the environment')
    root = Path(args.out)
    root.mkdir(parents=True, exist_ok=True, mode=0o700)
    with open(root / '.lock', 'w') as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        for task_id, payload in tasks:
            folder = root / task_id
            folder.mkdir(exist_ok=True, mode=0o700)
            write_atomic(folder / 'request.json', json.dumps(payload, ensure_ascii=False).encode())
        with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
            results = pool.map(lambda t: generate(root / t[0], t[1], api_key), tasks)
            for result in results:
                print(json.dumps(result))

if __name__ == '__main__':
    main()