from google import genai
from google.genai import types
from PIL import Image
import os
from dotenv import load_dotenv
from typing import List, Optional

load_dotenv()

client = genai.Client(api_key=os.getenv("BANANA_PRO_KEY"))

print(f"BANANA_PRO_KEY: {os.getenv('BANANA_PRO_KEY')}")

def generate_image_from_prompt_and_images(
    prompt: str,
    image_paths: List[str],
    aspect_ratio: str = "16:9",
    resolution: str = "2K"
) -> Optional[Image.Image]:
    """
    使用 Gemini API 生成图像
    
    参数:
        prompt: 文本提示词
        image_paths: 参考图像路径列表
        aspect_ratio: 图像宽高比，可选值: "1:1","2:3","3:2","3:4","4:3","4:5","5:4","9:16","16:9","21:9"
        resolution: 图像分辨率，可选值: "1K", "2K", "4K"
    
    返回:
        生成的图像对象 (PIL Image)，如果生成失败则返回 None
    """
    # 构建内容列表：包含提示词和所有参考图像
    contents = [prompt]
    for image_path in image_paths:
        contents.append(Image.open(image_path))
    
    # 调用 API 生成内容
    response = client.models.generate_content(
        model="gemini-3-pro-image-preview",
        contents=contents,
        config=types.GenerateContentConfig(
            response_modalities=['TEXT', 'IMAGE'],
            image_config=types.ImageConfig(
                aspect_ratio=aspect_ratio,
                image_size=resolution
            ),
        )
    )
    
    # 处理响应，提取图像
    for part in response.parts:
        if part.text is not None:
            print(part.text)
        elif image := part.as_image():
            return image
    
    return None


if __name__ == "__main__":
    # 测试代码
    prompt = "男人在月球上行走"
    image_paths = ["./test_output/output/front_portrait_0.png"]
    aspect_ratio = "16:9"
    resolution = "1K"
    
    generated_image = generate_image_from_prompt_and_images(
        prompt=prompt,
        image_paths=image_paths,
        aspect_ratio=aspect_ratio,
        resolution=resolution
    )
    
    if generated_image:
        generated_image.save("xxxxxxx.png")
        print("图像已保存到 xxxxxxx.png")