#!/usr/bin/env python3
"""Python 3.10+, standard library. Submit once, resume, save ALL returned images."""
import argparse
import base64
import json
import mimetypes
import os
import re
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

ORIGIN = 'https://openpix.cc'
LIMIT = 64 * 1024 * 1024  # Client download bound, not the API specification.


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


class APIError(RuntimeError):
    def __init__(self, code, message, retry_after=None):
        super().__init__(f'HTTP {code}: {message}')
        self.code, self.retry_after = code, retry_after


def request(key, path, payload=None, timeout=30):
    if not path.startswith('/v1/'):
        raise ValueError('Expected a fixed /v1/ path')
    req = urllib.request.Request(ORIGIN + path,
        data=None if payload is None else json.dumps(payload, ensure_ascii=False).encode('utf8'),
        headers={'Authorization': 'Bearer ' + key, 'Content-Type': 'application/json'},
        method='GET' if payload is None else 'POST')
    try:
        with urllib.request.build_opener(NoRedirect()).open(req, timeout=timeout) as response:
            result = json.load(response)
    except urllib.error.HTTPError as error:
        raise APIError(error.code, error.read(2048).decode('utf8', 'replace'), error.headers.get('Retry-After')) from error
    if not isinstance(result, dict) or (result.get('error') and not (result.get('status') or result.get('state'))):
        raise RuntimeError('API returned an error or invalid JSON object')
    return result


def prepare(body, references=(), force_edit=False):
    if not isinstance(body, dict) or not isinstance(body.get('model'), str):
        raise ValueError('Request must contain model')
    if not isinstance(body.get('prompt'), str) or not body['prompt'].strip():
        raise ValueError('prompt must be a nonempty string')
    body = dict(body)
    string_refs = body['model'].startswith(('mj-', 'seedream-'))
    if references:
        images = list(body.get('images', []))
        for name in references:
            file = Path(name)
            mime = mimetypes.guess_type(file.name)[0]
            if mime not in ('image/png', 'image/jpeg', 'image/webp'):
                raise ValueError('Local references must be PNG/JPEG/WebP')
            if not 0 < file.stat().st_size <= 25 * 1024 * 1024:
                raise ValueError('Example accepts each reference up to 25 MiB; model limits may be lower')
            url = 'data:' + mime + ';base64,' + base64.b64encode(file.read_bytes()).decode('ascii')
            images.append(url if string_refs else {'image_url': url})
        body['images'] = images
    edit = force_edit or (bool(body.get('images') or body.get('image')) and not string_refs)
    return body, '/v1/images/' + ('edits' if edit else 'generations') + '/async'


def task_id(value):
    if not isinstance(value, str) or not re.fullmatch(r'[A-Za-z0-9_-]+', value):
        raise ValueError('Invalid task ID; check logs before another POST')
    return value


def wait(key, identifier, seconds=1800):
    identifier = task_id(identifier)
    deadline, delay = time.monotonic() + seconds, 3
    while time.monotonic() < deadline:
        try:
            task = request(key, '/v1/images/tasks/' + identifier)
        except APIError as error:
            if error.code not in (429, 500, 502, 503, 504):
                raise
            try:
                delay = min(30, max(1, float(error.retry_after)))
            except (ValueError, TypeError):
                delay = min(30, delay * 1.5)
            time.sleep(delay)
            continue
        except (urllib.error.URLError, TimeoutError):
            time.sleep(delay)
            delay = min(30, delay * 1.5)
            continue
        returned = task.get('task_id') or task.get('id')
        if returned and returned != identifier:
            raise RuntimeError('Task ID mismatch')
        if task.get('submission_uncertain'):
            raise RuntimeError('Submission uncertain; inspect original task and billing, do not resubmit')
        status = str(task.get('status') or task.get('state') or '').lower()
        if status in ('completed', 'succeeded', 'success'):
            nested = task.get('result')
            data = task.get('data') or (nested.get('data') if isinstance(nested, dict) else None)
            if not isinstance(data, list) or not data or not all(isinstance(x, dict) for x in data):
                raise RuntimeError('Completed task has no valid images')
            return data, task
        if status in ('failed', 'failure', 'error', 'cancelled', 'canceled', 'expired'):
            raise RuntimeError('Task failed: ' + json.dumps(task.get('error'), ensure_ascii=False))
        if status not in ('processing', 'queued', 'running', 'pending', 'in_progress'):
            raise RuntimeError('Unknown task status: ' + status)
        value = task.get('poll_after_ms', 3000)
        delay = min(30, max(1, value / 1000)) if isinstance(value, (int, float)) else 3
        time.sleep(delay)
    raise TimeoutError('Resume using --task ' + identifier + '. Server task is not cancelled.')


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'
    if data[:6] in (b'GIF87a', b'GIF89a'): return '.gif'
    raise RuntimeError('Invalid or unsupported image signature; no output file replaced')


def save_images(items, directory):
    directory = Path(directory)
    directory.mkdir(parents=True, exist_ok=True)
    saved = []
    for i, item in enumerate(items, 1):
        if item.get('url'):
            url = urllib.parse.urlsplit(item['url'])
            if url.scheme not in ('https', 'http') or url.username or url.password or not url.hostname:
                raise ValueError('Invalid image result URL')
            # Separate request: no API Key sent to result URLs or redirects.
            with urllib.request.urlopen(item['url'], timeout=120) as response:
                if not response.headers.get('Content-Type', '').lower().startswith('image/'):
                    raise RuntimeError('Result URL did not return an image')
                data = response.read(LIMIT + 1)
        elif item.get('b64_json'):
            if len(item['b64_json']) > (LIMIT + 2) // 3 * 4:
                raise RuntimeError('Image exceeds client limit')
            data = base64.b64decode(item['b64_json'], validate=True)
        else:
            raise RuntimeError('Missing image URL/Base64')
        if not data or len(data) > LIMIT:
            raise RuntimeError('Image empty or exceeds client limit')
        target = directory / f'image-{i:02d}{image_extension(data)}'
        temp = target.with_suffix(target.suffix + '.part')
        temp.write_bytes(data)
        os.replace(temp, target)
        saved.append(str(target))
    return saved


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--request', help='UTF-8 JSON request; preserves model-specific fields')
    parser.add_argument('--model', default='gpt-image-2.5')
    parser.add_argument('--prompt')
    parser.add_argument('--size', default='1024x1024')
    parser.add_argument('--reference', action='append', default=[])
    parser.add_argument('--edit', action='store_true')
    parser.add_argument('--task', help='Resume; sends no generation POST')
    parser.add_argument('--timeout', type=int, default=1800)
    parser.add_argument('--output-dir', default='images')
    args = parser.parse_args()
    key = os.environ.get('MXB_API_KEY')
    if not key: parser.error('Set MXB_API_KEY')
    if args.timeout <= 0: parser.error('--timeout must be positive')
    directory = Path(args.output_dir)
    directory.mkdir(parents=True, exist_ok=True)
    identifier = args.task
    if not identifier:
        body = json.loads(Path(args.request).read_text(encoding='utf-8-sig')) if args.request else {
            'model': args.model, 'prompt': args.prompt, 'size': args.size, 'n': 1}
        body, path = prepare(body, args.reference, args.edit)
        state_file = directory / 'submission.json'
        state = {'state': 'submitting', 'created_at': int(time.time()), 'model': body['model'], 'path': path}
        # Exclusive state file prevents accidentally repeating a billable POST in the same directory.
        with state_file.open('x', encoding='utf8') as stream:
            json.dump(state, stream, ensure_ascii=False, indent=2)
        try:
            created = request(key, path, body, timeout=60)  # Exactly one POST.
            identifier = task_id(created.get('task_id') or created.get('id'))
        except Exception:
            state['state'] = 'submission_uncertain'
            state_file.write_text(json.dumps(state, indent=2), encoding='utf8')
            raise RuntimeError('Submission not retried. Inspect calling logs before creating another task.')
        state.update(state='accepted', task_id=identifier)
        state_file.write_text(json.dumps(state, indent=2), encoding='utf8')
        print('Task:', identifier, flush=True)
    data, final = wait(key, identifier, args.timeout)
    (directory / 'result.json').write_text(json.dumps(final, ensure_ascii=False, indent=2), encoding='utf8')
    for name in save_images(data, directory): print('Saved:', name)


if __name__ == '__main__': main()