import os
import time
import argparse
import requests
from dotenv import load_dotenv
from openai import OpenAI

# 加载根目录或当前目录的 .env 文件
env_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), ".env")
if os.path.exists(env_path):
    load_dotenv(env_path)
else:
    load_dotenv()

APIYI_KEY = os.getenv("APIYI_KEY")
if APIYI_KEY and not APIYI_KEY.startswith("sk-"):
    APIYI_KEY = f"sk-{APIYI_KEY}"
APIYI_BASE_URL = os.getenv("APIYI_BASE_URL", "https://api.apiyi.com/v1")

if not APIYI_KEY:
    print("错误: 请在 .env 文件中设置 APIYI_KEY")
    exit(1)

# 初始化客户端
client = OpenAI(
    api_key=APIYI_KEY,
    base_url=APIYI_BASE_URL
)

import tempfile
from urllib.parse import urlparse

def download_image_from_url(url):
    """从 URL 下载图片到临时文件"""
    try:
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        
        parsed_url = urlparse(url)
        extension = parsed_url.path.split('.')[-1] if '.' in parsed_url.path else 'jpg'
        
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=f'.{extension}')
        temp_file.write(response.content)
        temp_file.close()
        
        return temp_file.name
    except Exception as e:
        print(f"下载图片失败: {e}")
        return None

def generate_flux_image(prompt, aspect_ratio="1:1", model="flux-kontext-pro",
                       seed=None, safety_tolerance=2, output_format="jpeg",
                       prompt_upsampling=False, image_url=None, mask_url=None):
    """
    生成或编辑 Flux 图像
    """
    mode = "编辑" if image_url else "生成"
    print(f"正在{mode}图像...")
    print(f" - 模型: {model}")
    print(f" - 比例: {aspect_ratio}")
    print(f" - 提示词: {prompt}")
    
    image_path = None
    mask_path = None
    image_file = None
    mask_file = None
    
    try:
        # 构建 extra_body 参数
        extra_params = {
            "aspect_ratio": aspect_ratio,
            "safety_tolerance": safety_tolerance,
            "output_format": output_format
        }
        
        if seed is not None:
            extra_params["seed"] = seed
            
        if not image_url:
            # 纯生成特有参数
            extra_params["prompt_upsampling"] = prompt_upsampling
            response = client.images.generate(
                model=model,
                prompt=prompt,
                extra_body=extra_params
            )
        else:
            # 图生图编辑模式
            print(f" - 参考图: {image_url}")
            image_path = download_image_from_url(image_url)
            if not image_path:
                raise ValueError("无法下载参考图")
            image_file = open(image_path, "rb")
            
            edit_kwargs = {
                "model": model,
                "image": image_file,
                "prompt": prompt,
                "extra_body": extra_params
            }
            
            if mask_url:
                print(f" - 蒙版图: {mask_url}")
                mask_path = download_image_from_url(mask_url)
                if mask_path:
                    mask_file = open(mask_path, "rb")
                    edit_kwargs["mask"] = mask_file
                    
            response = client.images.edit(**edit_kwargs)
        
        image_out_url = response.data[0].url
        return image_out_url
    
    except Exception as e:
        print(f" 请求失败: {e}")
        return None
    finally:
        # 关闭文件并删除临时文件
        if image_file:
            image_file.close()
        if mask_file:
            mask_file.close()
        for path in [image_path, mask_path]:
            if path and os.path.exists(path):
                try:
                    os.unlink(path)
                except:
                    pass

def save_image_from_url(url, prompt, ratio):
    """从 URL 下载并保存图像"""
    if not url:
        return None
    try:
        print("正在下载临时链接中的图像...")
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        
        # 生成文件名，替换掉可能导致问题的非法字符
        timestamp = int(time.time())
        safe_prompt = ''.join(c if c.isalnum() else '_' for c in prompt[:30]).strip('_')
        
        # 保存到当前目录下的 outputs 文件夹
        output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "outputs")
        os.makedirs(output_dir, exist_ok=True)
        
        filename = f"flux_{safe_prompt}_{ratio.replace(':', 'x')}_{timestamp}.png"
        filepath = os.path.join(output_dir, filename)
        
        with open(filepath, 'wb') as f:
            f.write(response.content)
        
        print(f"✅ 已成功保存至: {filepath}")
        return filepath
        
    except requests.exceptions.RequestException as e:
        print(f"❌ 下载失败: {e}")
        print("提示：Flux URL 仅 10 分钟有效，请确保及时手动下载！")
        return None

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Flux API 图像生成高级调试工具")
    parser.add_argument("--prompt", type=str, default="A futuristic city with flying cars and neon lights, high resolution, cyberpunk style", help="图像描述提示词")
    parser.add_argument("--ratio", type=str, default="16:9", help="宽高比 (如: 1:1, 16:9, 9:16, 2:3, 3:2)")
    parser.add_argument("--model", type=str, default="flux-kontext-pro", choices=["flux-kontext-pro", "flux-kontext-max"], help="模型名称")
    parser.add_argument("--seed", type=int, default=None, help="随机种子 (默认随机)")
    parser.add_argument("--format", type=str, default="png", choices=["jpeg", "png"], help="输出图片格式")
    parser.add_argument("--upsampling", action="store_true", help="开启提示词自动增强")
    
    args = parser.parse_args()
    
    image_url = generate_flux_image(
        prompt=args.prompt,
        aspect_ratio=args.ratio,
        model=args.model,
        seed=args.seed,
        output_format=args.format,
        prompt_upsampling=args.upsampling
    )
    
    if image_url:
        print(f"\n✅ 生成成功！图像 URL (10分钟有效): \n{image_url}")
        save_image_from_url(image_url, args.prompt, args.ratio)