#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
🍌 Nano Banana (千层套路·多模态生成) 外部调用脚本

【工具定位】：它使用的是 Google 最新的多模态 Imagen 3 引擎（gemini-3.1-flash-image-preview）。
【独门绝技 - 跨模态多图推理】：
与传统的 SD / Flux (图生图) 的本质区别在于，你可以给它丢【任意张完全不同的图片】，
然后在 prompt 里随心所欲地让它“融梗”交汇。
比如传入图片 [猫.jpg] 和 [未来城.jpg]，让它“把这只猫生成在这座未来城里”。

【参数模式支持】：
1. 纯文生图：只传 prompt
2. 单图生图/重绘：传 1 张图 + prompt
3. 多路意象融合：传 N 张图 + prompt

【本地文件自动上传机制】：
大模型 API 往往更喜欢吃稳定的 CDN 外链。
此脚本在启动前会自动检测你传入的图片：如果是本地硬盘文件（比如 examples/cat.png），
脚本会自动静默调用内部的 OSS 工具，把它秒传为 https://res.cybertogether.net/.. 干净外链，然后投喂给大脑！
"""

import asyncio
import os
import argparse
import sys
import httpx
import json

# 动态引入我们系统现成的 CDN 上传脚本
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'examples', 'production_restore'))
try:
    from upload import upload_image
except Exception as e:
    import traceback
    print(f"错误: 导入 upload.py 及其依赖链失败: {e}\n{traceback.format_exc()}")
    sys.exit(1)

ROUTER_URL = "http://43.106.118.91:8001/run_tool"

async def process_images(images_list: list[str]) -> list[str]:
    """处理图片数组，外链直接过，本地文件传 OSS"""
    final_urls = []
    for item in images_list:
        if item.startswith("http://") or item.startswith("https://"):
            print(f"✅ 检测到公网外链，无需转存: {item}")
            final_urls.append(item)
        elif os.path.exists(item):
            print(f"📦 正在极速倒卖本地图片到 CDN: {item}")
            try:
                uploaded_url = await upload_image(item, 'aigc-admin', 'crawler/image')
                if uploaded_url:
                    print(f"🚀 上传成功: {uploaded_url}")
                    final_urls.append(uploaded_url)
                else:
                    print(f"❌ 上传失败: {item}")
            except Exception as e:
                print(f"❌ 上传报错: {e}")
        else:
            print(f"⚠️ 跳过找不到的本地文件或无法识别的格式: {item}")
    return final_urls

async def run_nano_banana(prompt: str, images: list[str] = None, model: str = None, aspect_ratio: str = None):
    print(f"\n=======================")
    print(f"🍌 Nano Banana 启动中...")
    print(f"=======================")

    # 1. 整理图片
    final_image_urls = []
    if images and len(images) > 0:
        print(f"🔍 检查到传了 {len(images)} 张神秘原图，准备过安检...")
        final_image_urls = await process_images(images)
    
    # 2. 组装发给大模型中枢 Router 的参数
    params = {
        "prompt": prompt,
        "image_urls": final_image_urls if final_image_urls else None
    }
    
    if model:
        params["model"] = model
        
    if aspect_ratio:
        params["aspect_ratio"] = aspect_ratio

    payload = {
        "tool_id": "nano_banana",
        "params": params
    }

    # 3. 轰入 API
    print("\n⚡ 正在呼叫总后台路由节点打怪...")
    try:
        async with httpx.AsyncClient(timeout=300.0) as client:
            resp = await client.post(ROUTER_URL, json=payload)
            resp.raise_for_status()
            
            result = resp.json()
            if result.get("status") == "success":
                gen_data = result.get("result", {})
                
                print("\n🎉 === 生成成功！ ===")
                # 打印文本回复
                if gen_data.get("text"):
                    print(f"\n💬 模型回话:\n{gen_data.get('text')}")
                
                # 打印图片输出
                images_out = gen_data.get("images", [])
                if images_out:
                    print("\n🖼️ 吐出的神图 (Base64 数据流已转为你本地文件)：")
                    for idx, img_b64 in enumerate(images_out):
                        # 处理前缀 data:image/jpeg;base64,
                        if img_b64.startswith("data:"):
                            mime_split = img_b64.split(";base64,")
                            if len(mime_split) == 2:
                                ext = mime_split[0].split("/")[-1]
                                raw_data = mime_split[1]
                                
                                import base64
                                save_path = f"banana_output_{idx}.{ext}"
                                with open(save_path, "wb") as f:
                                    f.write(base64.b64decode(raw_data))
                                print(f" 💾 已保存到本地 -> {save_path}")
            else:
                print(f"❌ 大模型傲娇了: {result.get('error')}")
                
    except Exception as e:
        print(f"💥 网络大爆炸报错: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="调用 Nano Banana 进行任意风格的多段融合魔法")
    parser.add_argument("-p", "--prompt", type=str, required=True, help="你想对 AI 喊瞎什么 (比如：用图1的赛博风画一只图2里的猫)")
    parser.add_argument("-i", "--images", type=str, nargs="+", help="无限追加的垫图清单（可以是现成的 http 链接，也可以是你电脑里的硬盘文件如 example.png）")
    parser.add_argument("-m", "--model", type=str, default=None, help="覆盖模型 (默认后台会走 gemini-3.1-flash-image-preview)")
    parser.add_argument("-a", "--aspect_ratio", type=str, default=None, help="图片比例，例如 3:4, 16:9, 1:1 等")
    
    args = parser.parse_args()
    
    asyncio.run(run_nano_banana(prompt=args.prompt, images=args.images, model=args.model, aspect_ratio=args.aspect_ratio))