import requests, json, base64, os, sys

# ============ 配置区 ============
# ⚠️ 生图慢，timeout≥300s
CONFIG_KEY = '<YOUR_CONFIG_KEY>'  # mykey.py中的配置变量名，需指向支持GPT-image-1的配置（如gpt5.4/5.5）
# ================================

_DIR = os.path.dirname(os.path.abspath(__file__))
for _p in [os.path.join(_DIR, '..'), os.path.join(_DIR, '../..')]:
    if _p not in sys.path: sys.path.insert(0, _p)

def generate_image(prompt, save_path=None, size="auto", quality="high", timeout=180):
    """GPT图像生成。返回保存路径，失败返回 'Error: ...'
    
    用法: path = generate_image("一只猫骑自行车", "cat.png")
    注意: 必须用stream模式，中转站非流式会剥离图片数据
    """
    import mykey
    cfg = getattr(mykey, CONFIG_KEY)
    if not save_path:
        save_path = os.path.join(os.getcwd(), 'generated.png')

    url = cfg['apibase'].rstrip('/') + '/responses'
    headers = {'Authorization': f'Bearer {cfg["apikey"]}', 'Content-Type': 'application/json'}
    payload = {
        "model": cfg['model'], "stream": True,
        "input": prompt,
        "tools": [{"type": "image_generation", "size": size, "quality": quality}],
    }

    try:
        resp = requests.post(url, headers=headers, json=payload, stream=True, timeout=timeout)
        resp.raise_for_status()
        # iter_content收集完整响应，避免iter_lines截断超大base64行
        buf = b''
        for chunk in resp.iter_content(chunk_size=1024*64):
            buf += chunk
        image_b64 = None
        for line in buf.decode('utf-8', errors='replace').split('\n'):
            if not line.startswith('data: '): continue
            raw = line[6:]
            if raw == '[DONE]': break
            try:
                evt = json.loads(raw)
            except json.JSONDecodeError:
                continue
            item = evt.get('item', {})
            if item.get('type') == 'image_generation_call' and item.get('result'):
                image_b64 = item['result']
                break
        if not image_b64:
            return "Error: 未获取到图片数据"
        img_bytes = base64.b64decode(image_b64)
        with open(save_path, 'wb') as f:
            f.write(img_bytes)
        print(f"  🖼️ 已保存: {save_path} ({len(img_bytes)//1024}KB)")
        return save_path
    except requests.exceptions.Timeout:
        return f"Error: 请求超时 (>{timeout}s)"
    except Exception as e:
        return f"Error: {type(e).__name__}: {e}"

if __name__ == '__main__':
    print('✅ Image Gen API 已就绪 (GPT图像生成，stream模式)')