"""Python 3.8+: pip install requests pillow. Native Gemini image generation/editing."""
import argparse
import base64
import io
import json
import os
from pathlib import Path
from urllib.parse import quote

import requests
from PIL import Image


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--prompt', required=True)
    parser.add_argument('--image', action='append', default=[])
    parser.add_argument('--size', choices=['1K', '2K', '4K'], default='1K')
    parser.add_argument('--aspect', default='1:1')
    parser.add_argument('--history', help='Private JSON file storing full native contents for subsequent edits')
    parser.add_argument('--out', default='gemini-images-output')
    args = parser.parse_args()
    key, model = os.environ['YUNQIAI_API_KEY'], os.environ['YUNQIAI_MODEL']
    if model not in ['gemini-3.1-flash-image', 'gemini-3-pro-image-preview', 'gemini-nano-banana-2.1']:
        raise ValueError('Choose an exact Gemini image model ID')
    history = []
    history_path = Path(args.history) if args.history else None
    if history_path and history_path.exists():
        saved = json.loads(history_path.read_text(encoding='utf-8'))
        if saved.get('model') != model:
            raise ValueError('History model differs from selected model')
        history = saved['contents']
    parts = [{'text': args.prompt}]
    for name in args.image:
        path = Path(name)
        with Image.open(path) as image:
            image.verify()
            mime = {'PNG': 'image/png', 'JPEG': 'image/jpeg', 'WEBP': 'image/webp'}.get(image.format)
        if not mime:
            raise ValueError('Reference must be PNG, JPEG or WebP')
        parts.append({'inlineData': {'mimeType': mime, 'data': base64.b64encode(path.read_bytes()).decode('ascii')}})
    history.append({'role': 'user', 'parts': parts})
    body = {'contents': history, 'generationConfig': {'maxOutputTokens': 8192, 'responseModalities': ['TEXT', 'IMAGE'], 'imageConfig': {'aspectRatio': args.aspect, 'imageSize': args.size}}}
    response = requests.post('https://www.yunqiai.chat/v1beta/models/' + quote(model, safe='') + ':generateContent', headers={'x-goog-api-key': key}, json=body, allow_redirects=False, timeout=(15, 300))
    if not 200 <= response.status_code < 300:
        raise RuntimeError('HTTP ' + str(response.status_code) + ': ' + response.text[:4000])
    data = response.json()
    if data.get('error'):
        raise RuntimeError(json.dumps(data['error'], ensure_ascii=False))
    candidates = data.get('candidates', [])
    if not candidates or candidates[0].get('finishReason') != 'STOP':
        raise RuntimeError('No completed image candidate: ' + str([c.get('finishReason') for c in candidates]))
    content = candidates[0]['content']
    if history_path:
        history.append(content)
        history_path.parent.mkdir(parents=True, exist_ok=True)
        history_path.write_text(json.dumps({'model': model, 'contents': history}, ensure_ascii=False), encoding='utf-8')
    output = Path(args.out)
    output.mkdir(parents=True, exist_ok=True)
    saved, errors = [], []
    for index, part in enumerate(content.get('parts', [])):
        if part.get('thought'):
            continue
        inline = part.get('inlineData') or part.get('inline_data')
        if not inline:
            if part.get('text'):
                print(part['text'])
            continue
        try:
            raw = base64.b64decode(inline['data'], validate=True)
            with Image.open(io.BytesIO(raw)) as image:
                image.load()
                extension = {'PNG': 'png', 'JPEG': 'jpg', 'WEBP': 'webp'}[image.format]
                target = output / ('image-' + str(index) + '.' + extension)
                target.write_bytes(raw)
                saved.append({'file': str(target), 'size': list(image.size)})
                if max(image.size) < {'1K': 1000, '2K': 2000, '4K': 4000}[args.size]:
                    errors.append('Final image smaller than requested size tier: ' + str(target))
        except Exception as error:
            errors.append(str(error))
    print(json.dumps({'images': saved, 'errors': errors}, ensure_ascii=False))
    if not saved or errors:
        raise RuntimeError('Missing or invalid final images; do not automatically submit again')


if __name__ == '__main__':
    main()