#!/usr/bin/env python3
"""
AI 图片/视频生成工具
支持平台: 通义万相(免费额度)、海螺AI、MiniMax API
用法:
  python3 ai_gen.py image "一只橘猫在沙发上" [--model wax3]
  python3 ai_gen.py video "海浪拍打沙滩" [--model minimax-video-01]
  python3 ai_gen.py status
"""

import os, sys, json, time, argparse, subprocess, urllib.parse, urllib.request
from pathlib import Path

# ─── 配置区 ──────────────────────────────────────────
# 通义万相（阿里云百炼）- 有免费额度
# 获取: https://bailian.console.aliyun.com/
DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY", "")
# 体验地址: https://tongyi.aliyun.com/wanxiang/

# 海螺AI (MiniMax) - API
# 获取: https://www.minimax.io/ 或 https://platform.minimaxi.com/
MINIMAX_API_KEY = os.environ.get("MINIMAX_API_KEY", "")
MINIMAX_GROUP_ID = os.environ.get("MINIMAX_GROUP_ID", "")

# ─── 工具函数 ─────────────────────────────────────────

def run(cmd):
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    return r.stdout.strip(), r.stderr.strip(), r.returncode

def save_image(url, name=None):
    if not name:
        name = f"/tmp/ai_gen_{int(time.time())}.png"
    subprocess.run(f"curl -sL '{url}' -o '{name}'", shell=True)
    print(f"  📷 已保存: {name}")
    return name

def save_video(url, name=None):
    if not name:
        name = f"/tmp/ai_gen_{int(time.time())}.mp4"
    subprocess.run(f"curl -sL '{url}' -o '{name}'", shell=True)
    print(f"  🎬 已保存: {name}")
    return name

# ─── 通义万相 ──────────────────────────────────────────
import requests

def wanxiang_image(prompt, model="wanx3.1-t2i-turbo"):
    """通义万相 - 图像生成（免费额度）"""
    if not DASHSCOPE_API_KEY:
        print("  ⚠️  未设置 DASHSCOPE_API_KEY")
        print("     体验地址: https://tongyi.aliyun.com/wanxiang/")
        print("     或获取API: https://bailian.console.aliyun.com/")
        return None

    url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"
    headers = {
        "Authorization": f"Bearer {DASHSCOPE_API_KEY}",
        "Content-Type": "application/json",
    }
    data = {
        "model": model,
        "input": {"prompt": prompt},
        "parameters": {
            "size": "1024*1024",
            "n": 1,
            "style": "<auto>",
        }
    }
    resp = requests.post(url, headers=headers, json=data, timeout=30)
    result = resp.json()
    if resp.status_code != 200:
        print(f"  ❌ 错误 {resp.status_code}: {result}")
        return None

    task_id = result.get("output", {}).get("task_id")
    if not task_id:
        print(f"  ❌ 无task_id: {result}")
        return None

    # 轮询结果
    print(f"  ⏳ 任务ID: {task_id}，等待生成...")
    for i in range(30):
        time.sleep(2)
        status_url = f"https://dashscope.aliyuncs.com/api/v1/tasks/{task_id}"
        s = requests.get(status_url, headers=headers, timeout=10).json()
        status = s.get("output", {}).get("task_status", "")
        print(f"    [{i+1}] {status}")
        if status == "succeeded":
            img_url = s["output"]["results"][0]["image_url"]
            return save_image(img_url)
        elif status == "failed":
            print(f"  ❌ 生成失败: {s}")
            return None
    print("  ❌ 超时")
    return None

# ─── Pollinations（免费，无需API Key）──────────────────

def pollinations_image(prompt, width=1024, height=1024, model=None):
    """Pollinations - 完全免费的AI生图"""
    encoded_prompt = urllib.parse.quote(prompt)
    url = f"https://image.pollinations.ai/prompt/{encoded_prompt}?width={width}&height={height}"
    if model:
        url += f"&model={model}"
    save_path = f"/tmp/pollen_{int(time.time())}.png"
    print(f"  🔗 URL: {url}")
    print(f"  ⏳ 生成中（通常10-30秒）...")
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"})
        with urllib.request.urlopen(req, timeout=60) as resp:
            with open(save_path, "wb") as f:
                f.write(resp.read())
        size = os.path.getsize(save_path)
        print(f"  ✅ 已保存: {save_path} ({size//1024}KB)")
        return save_path
    except Exception as e:
        print(f"  ❌ 下载失败: {e}")
        return None

def pollinations_video(prompt, duration=5):
    """Pollinations - 免费视频生成（experimental）"""
    encoded_prompt = urllib.parse.quote(prompt)
    url = f"https://video.pollinations.ai/v1?prompt={encoded_prompt}&duration={duration}"
    save_path = f"/tmp/pollen_vid_{int(time.time())}.mp4"
    print(f"  🔗 {url}")
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"})
        with urllib.request.urlopen(req, timeout=60) as resp:
            with open(save_path, "wb") as f:
                f.write(resp.read())
        size = os.path.getsize(save_path)
        print(f"  ✅ 已保存: {save_path} ({size//1024}KB)")
        return save_path
    except Exception as e:
        print(f"  ❌ 下载失败: {e}")
        return None

# ─── 海螺AI (MiniMax) ─────────────────────────────────

def hailuo_image(prompt, model="image-01"):
    """海螺AI - 图像生成"""
    if not MINIMAX_API_KEY:
        print("  ⚠️  未设置 MINIMAX_API_KEY")
        print("     获取: https://platform.minimaxi.com/")
        return None

    url = "https://api.minimax.chat/v1/image_generation"
    headers = {
        "Authorization": f"Bearer {MINIMAX_API_KEY}",
        "Content-Type": "application/json",
    }
    data = {
        "model": model,
        "prompt": prompt,
        "aspect_ratio": "1:1",
        "response_format": "url",
    }
    resp = requests.post(url, headers=headers, json=data, timeout=30)
    result = resp.json()
    if "data" in result and result["data"]:
        img_url = result["data"][0]["url"]
        return save_image(img_url)
    else:
        print(f"  ❌ 错误: {result}")
        return None

def hailuo_video(prompt, model="video-01"):
    """海螺AI - 视频生成（推荐，免费额度高）"""
    if not MINIMAX_API_KEY or not MINIMAX_GROUP_ID:
        print("  ⚠️  未设置 MINIMAX_API_KEY 或 MINIMAX_GROUP_ID")
        print("     获取: https://platform.minimaxi.com/")
        return None

    url = "https://api.minimax.chat/v1/video_generation"
    headers = {
        "Authorization": f"Bearer {MINIMAX_API_KEY}",
        "Content-Type": "application/json",
        "GroupId": MINIMAX_GROUP_ID,
    }
    data = {
        "model": model,
        "prompt": prompt,
        "duration": 5,
        "resolution": "720p",
        "response_format": "url",
    }
    resp = requests.post(url, headers=headers, json=data, timeout=30)
    result = resp.json()
    print(f"  📋 响应: {json.dumps(result, ensure_ascii=False)[:200]}")

    task_id = result.get("data", {}).get("task_id") or result.get("task_id")
    if not task_id:
        print(f"  ❌ 无task_id: {result}")
        return None

    print(f"  ⏳ 任务ID: {task_id}，等待生成（视频约需30-60秒）...")
    for i in range(60):
        time.sleep(3)
        # 查询状态
        poll_url = f"https://api.minimax.chat/v1/video_generation?task_id={task_id}"
        s = requests.get(poll_url, headers=headers, timeout=10).json()
        status = s.get("data", {}).get("status", "processing")
        print(f"    [{i+1}] {status}")
        if status == "success" or status == "completed":
            video_url = s["data"].get("video_info", {}).get("video_url", "")
            if not video_url:
                # 尝试其他字段
                video_url = s["data"].get("video_url", "")
            if video_url:
                return save_video(video_url)
            print(f"  📋 结果: {json.dumps(s, ensure_ascii=False)[:300]}")
            return None
        elif status == "failed":
            print(f"  ❌ 生成失败: {s}")
            return None
    print("  ❌ 超时")
    return None

# ─── 主程序 ───────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="AI 图片/视频生成工具")
    sub = parser.add_subparsers(dest="cmd")

    p_img = sub.add_parser("image", help="生成图片")
    p_img.add_argument("prompt", help="图片描述")
    p_img.add_argument("--model", default="pollen", help="模型: pollen(默认/免费), wanxiang, hailuo")
    p_img.add_argument("--width", type=int, default=1024)
    p_img.add_argument("--height", type=int, default=1024)

    p_vid = sub.add_parser("video", help="生成视频")
    p_vid.add_argument("prompt", help="视频描述")
    p_vid.add_argument("--model", default="hailuo", help="模型: hailuo (默认 hailuo)")

    sub.add_parser("status", help="查看配置状态")

    args = parser.parse_args(sys.argv[1:] if len(sys.argv) > 1 else ["status"])

    if args.cmd == "status":
        print("=" * 40)
        print("AI 生成工具状态")
        print("=" * 40)
        print(f"  Pollinations (图像): ✅ 随时可用（免费）")
        has_dash = bool(DASHSCOPE_API_KEY)
        has_mini = bool(MINIMAX_API_KEY)
        print(f"  通义万相 (图像): {'✅ 已配置' if has_dash else '❌ 未配置'}")
        print(f"  海螺AI (图像+视频): {'✅ 已配置' if has_mini else '❌ 未配置'}")
        print()
        print("获取API密钥:")
        print("  通义万相: https://bailian.console.aliyun.com/ (有免费额度)")
        print("  海螺AI:   https://platform.minimaxi.com/")
        print()
        print("使用方法:")
        print("  python3 ai_gen.py image '一只橘猫'                    # 默认 Pollinations")
        print("  python3 ai_gen.py image '一只橘猫' --width 1024 --height 1024")
        print("  python3 ai_gen.py image '风景' --model wanxiang       # 通义万相")
        print("  python3 ai_gen.py video '海浪拍打沙滩'                # 海螺AI视频")
        print("  python3 ai_gen.py status")
        return

    if args.cmd == "image":
        print(f"\n🎨 生成图片: {args.prompt}")
        if args.model == "pollen":
            print("  平台: Pollinations（免费，无需API Key）")
            result = pollinations_image(args.prompt, width=args.width, height=args.height)
        elif args.model == "wanxiang":
            print("  平台: 通义万相")
            result = wanxiang_image(args.prompt)
        elif args.model == "hailuo":
            print("  平台: 海螺AI")
            result = hailuo_image(args.prompt)
        else:
            print(f"  ❌ 未知模型: {args.model}")
            result = None
        if result:
            print(f"\n✅ 完成: {result}")

    elif args.cmd == "video":
        print(f"\n🎬 生成视频: {args.prompt}")
        result = hailuo_video(args.prompt)
        if result:
            print(f"\n✅ 完成: {result}")

    else:
        parser.print_help()

if __name__ == "__main__":
    main()