#!/usr/bin/env python3
"""Flux API quickstart. Standard library only. No key is written to disk."""
import argparse
import http.client
import json
import os
import shutil
import time
import urllib.error
import urllib.request
import uuid
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--check', action='store_true', help='Check connection and balance without generating')
    parser.add_argument('--state', default='flux-job.json', help='Use a new file for each NEW generation')
    parser.add_argument('--output-dir', help='Result folder (default: next to the state file)')
    parser.add_argument('--wait-seconds', type=int, default=600)
    args = parser.parse_args()
    base = os.environ.get('FLUX_API_URL', 'https://api.flux-context.org/api/v1').rstrip('/')
    key = os.environ['FLUX_API_KEY']

    def call(path, payload=None, idem=None):
        headers = {'Authorization': 'Bearer ' + key}
        if idem:
            headers['Idempotency-Key'] = idem
        data = None
        if payload is not None:
            headers['Content-Type'] = 'application/json'
            data = json.dumps(payload).encode()
        request = urllib.request.Request(base + path, data=data, headers=headers)
        try:
            with urllib.request.urlopen(request, timeout=75) as response:
                return response.status, json.load(response), 10
        except urllib.error.HTTPError as error:
            try:
                body = json.load(error)
            except (ValueError, UnicodeError):
                body = {'error': {'message': 'Unexpected server response. Retry the same job.'}}
            retry = error.headers.get('Retry-After', '10')
            return error.code, body, min(60, max(1, int(retry))) if retry.isdigit() else 10

    if args.check:
        try:
            status, body, _ = call('/balance')
        except (urllib.error.URLError, http.client.HTTPException, OSError, ValueError):
            print('Connection check failed. Check your network and try again.')
            return 1
        print(json.dumps(body, indent=2))
        return 0 if status == 200 else 1

    state_file = Path(args.state)
    if state_file.exists():
        job = json.loads(state_file.read_text())
    else:
        job = {'idempotency_key': str(uuid.uuid4()), 'payload': {
            'model': 'gpt-image-2',
            'input': {'prompt': 'Editorial fashion portrait', 'aspect_ratio': '4:5',
                      'resolution': '2k', 'quality': 'high', 'image_urls': []}}}
        # Persist BEFORE the POST. Restarts reuse the same key and body.
        with state_file.open('x') as file:
            json.dump(job, file)

    deadline = time.monotonic() + args.wait_seconds
    while time.monotonic() < deadline:
        try:
            if job.get('id'):
                status, body, delay = call('/generations/' + job['id'])
            else:
                status, body, delay = call('/generations', job['payload'], job['idempotency_key'])
        except (urllib.error.URLError, http.client.HTTPException, OSError, ValueError):
            time.sleep(min(10, max(0, deadline - time.monotonic())))
            continue  # Reuse the same ID or idempotency key; never create a new job.
        if body.get('id'):
            job['id'] = body['id']
            temp = state_file.with_suffix(state_file.suffix + '.tmp')
            temp.write_text(json.dumps(job))
            temp.replace(state_file)
        if body.get('status') in ('succeeded', 'failed', 'rejected'):
            print(json.dumps(body, indent=2))
            if body['status'] != 'succeeded':
                return 1
            output = Path(args.output_dir) if args.output_dir else state_file.parent / (state_file.stem + '-results')
            output.mkdir(parents=True, exist_ok=True)
            for index, result in enumerate(body.get('results', [])):
                if not result.get('url') or result.get('expired'):
                    print('Result retention has expired. Generation records remain available.')
                    return 1
                try:
                    # Signed result URLs authorize downloads; never forward the API key.
                    with urllib.request.urlopen(result['url'], timeout=75) as media:
                        kind = media.headers.get_content_type()
                        extension = {'image/png': 'png', 'image/jpeg': 'jpg', 'image/webp': 'webp', 'video/mp4': 'mp4'}.get(kind, 'bin')
                        target = output / (str(index + 1) + '.' + extension)
                        temporary = target.with_suffix(target.suffix + '.part')
                        with temporary.open('wb') as file:
                            shutil.copyfileobj(media, file)
                        temporary.replace(target)
                        print('Saved:', target)
                except (urllib.error.URLError, http.client.HTTPException, OSError, ValueError):
                    print('Download interrupted. Run again with the SAME --state file to refresh the link.')
                    return 2
            return 0
        if status not in (200, 202, 429, 503) and body.get('error', {}).get('code') != 'SUBMISSION_IN_PROGRESS':
            print(json.dumps(body, indent=2))
            return 1
        time.sleep(min(delay, max(0, deadline - time.monotonic())))
    print('Still pending. Run again with the SAME --state file. Request ID:', job.get('id', 'not confirmed'))
    return 2


if __name__ == '__main__':
    raise SystemExit(main())