"""本站 Gemini SSE 三轮改图；Python 3 标准库，最多生成三张 1K 图片。"""
import argparse
import base64
import getpass
import json
import os
from pathlib import Path
import uuid
from urllib.error import HTTPError
from urllib.request import Request, urlopen

BASE = 'https://yyds.seqnode.com'
MODELS = ('gemini-3.1-flash-image', 'gemini-nano-banana-2.1')
PROMPTS = (
    'A blue ceramic teacup on a pale wooden table, one orange beside it, soft daylight, no text.',
    'Edit the previous image: change only the blue teacup to bright red. Preserve the orange, table, lighting and composition.',
    'Edit the latest image: add one silver teaspoon to the right of the red cup. Preserve the red cup, orange and the rest of the scene.',
)

def sse_events(lines):
    pending = []
    total = 0
    for line in lines:
        total += len(line)
        if total > 128 * 1024 * 1024:
            raise ValueError('响应超过示例的 128 MiB 上限')
        if line.strip() == b'':
            if pending:
                data = b'\n'.join(pending)
                pending.clear()
                if data != b'[DONE]':
                    yield json.loads(data)
        elif line.startswith(b'data:'):
            pending.append(line[5:].lstrip(b' ').rstrip(b'\r\n'))
    if pending and b'\n'.join(pending) != b'[DONE]':
        yield json.loads(b'\n'.join(pending))

def collect(events):
    parts, finish, count = [], None, 0
    for event in events:
        count += 1
        if event.get('error'):
            raise RuntimeError('SSE 返回错误：' + str(event['error'].get('code', 'unknown')))
        for candidate in event.get('candidates', []):
            if candidate.get('index', 0) == 0:
                # ponytail: 按事件顺序保留原始 parts；不拼接、解码或重造签名。
                parts.extend(candidate.get('content', {}).get('parts', []))
                finish = candidate.get('finishReason', finish)
    if finish != 'STOP':
        raise RuntimeError('响应未正常完成，停止后续修改：' + str(finish))
    if not any((p.get('inlineData') or p.get('inline_data')) and not p.get('thought') for p in parts):
        raise RuntimeError('响应没有最终图片，停止后续修改')
    return {'role': 'model', 'parts': parts}, count

def run(model, key, turns=3):
    history = []
    output = Path('banana-stream-' + uuid.uuid4().hex[:10])
    output.mkdir()
    for turn, prompt in enumerate(PROMPTS[:turns], 1):
        history.append({'role': 'user', 'parts': [{'text': prompt}]})
        body = {'contents': history, 'generationConfig': {'responseModalities': ['IMAGE'], 'imageConfig': {'imageSize': '1K', 'aspectRatio': '1:1'}}}
        request = Request(BASE + '/v1beta/models/' + model + ':streamGenerateContent?alt=sse', data=json.dumps(body).encode(), headers={'Authorization': 'Bearer ' + key, 'Content-Type': 'application/json', 'Accept': 'text/event-stream'})
        try:
            with urlopen(request, timeout=300) as response:
                if 'text/event-stream' not in response.headers.get('Content-Type', ''):
                    raise RuntimeError('响应不是 SSE，请检查端点和模型')
                content, count = collect(sse_events(response))
        except HTTPError as exc:
            raise RuntimeError('HTTP ' + str(exc.code) + '；停止调用，请检查服务端请求日志') from None
        for index, part in enumerate(content['parts']):
            media = part.get('inlineData') or part.get('inline_data')
            if not media or part.get('thought'):
                continue
            mime = media.get('mimeType', media.get('mime_type'))
            suffix = {'image/jpeg': '.jpg', 'image/png': '.png', 'image/webp': '.webp'}.get(mime)
            if not suffix:
                raise RuntimeError('示例未支持的图片格式：' + str(mime))
            target = output / ('turn-' + str(turn) + '-' + str(index) + suffix)
            target.write_bytes(base64.b64decode(media['data'], validate=True))
            print('图片：', target.resolve())
        signatures = sum(bool(p.get('thoughtSignature') or p.get('thought_signature')) for p in content['parts'])
        print('轮次', turn, 'SSE 事件数', count, '签名数', signatures)
        if not signatures and turn < turns:
            raise RuntimeError('未返回签名，停止本示例的签名多轮测试')
        history.append(content)  # 下一轮回传完整模型内容，图片和签名保持原样。

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--model', choices=MODELS, default=MODELS[0])
    parser.add_argument('--turns', type=int, choices=(1, 2, 3), default=3)
    args = parser.parse_args()
    api_key = os.environ.get('YM2API_KEY') or getpass.getpass('本站 API Key（输入隐藏）：').strip()
    if not api_key:
        parser.error('请提供本站 API Key')
    run(args.model, api_key, args.turns)