#!/usr/bin/env python3
"""
Agnes 图片生成工具（多风格）
============================
基于 agnes-image-2.0-flash 模型，一句话出图。
内置 8 种风格预设，覆盖人物/国风/仙宫/现代/Zine。

用法:
  python3 agnes_image.py "一只猫在花园"                      # 默认风格
  python3 agnes_image.py "江南水乡" --style guofeng           # 国风建筑
  python3 agnes_image.py "仙女宫殿" --style xian              # 仙宫天宫
  python3 agnes_image.py "美术馆" --style modern              # 现代建筑
  python3 agnes_image.py "清冷美女" --style char              # 人物角色
  python3 agnes_image.py "禅意庭院" --style wabi              # 日式侘寂
  python3 agnes_image.py "赛博城市" --style cyber             # 赛博朋克
  python3 agnes_image.py "水墨山水" --style ink               # 水墨风格
  python3 agnes_image.py "描述" --style guofeng --size 1792x1024  # 自定义尺寸
  python3 agnes_image.py "描述" --output /tmp/my.png          # 指定输出
  python3 agnes_image.py "描述" --push                        # 生成后推飞书
  python3 agnes_image.py list                                  # 列出缓存
"""

import json, sys, os, time, subprocess
from datetime import datetime
from pathlib import Path

API_BASE = "https://apihub.agnes-ai.com/v1"
MODEL = "agnes-image-2.0-flash"
FALLBACK_MODEL = "agnes-image-2.1-flash"
OUTPUT_DIR = Path.home() / ".hermes" / "image_cache"
OUTPUT_DIR.mkdir(exist_ok=True)

FEISHU_WEBHOOK = "https://open.feishu.cn/open-apis/bot/v2/hook/446db983-e392-4d2c-bfb8-f9060e5df3ad"

# === 风格预设 ===
STYLES = {
    "default": {
        "name": "默认",
        "suffix": "",
        "negative": "low quality, blurry, watermark, logo, text",
    },
    "guofeng": {
        "name": "国风建筑",
        "suffix": (
            ", Chinese traditional architecture, flying eaves, blue-grey tiles, "
            "red lanterns, stone bridge, pavilion, misty mountains, "
            "cinematic lighting, environment concept art, game scene design, "
            "aerial view, wide angle, dramatic sunset, warm lantern lights, "
            "highly detailed, 8k, trending on artstation, epic scale"
        ),
        "negative": (
            "low quality, blurry, bad perspective, distorted architecture, "
            "modern elements, readable text, watermark, logo, broken buildings"
        ),
    },
    "xian": {
        "name": "仙宫天宫",
        "suffix": (
            ", xianxia celestial palace, white marble pillars, golden ornaments, "
            "sea of clouds, floating palace, ancient pine, full moon, "
            "ethereal mist, celestial atmosphere, dragon carvings, "
            "cinematic lighting, environment concept art, game scene design, "
            "ultra detailed, 8k, epic scale, mysterious"
        ),
        "negative": (
            "low quality, blurry, bad perspective, distorted architecture, "
            "modern elements, readable text, watermark, logo, "
            "floating without support, wrong scale"
        ),
    },
    "modern": {
        "name": "现代建筑",
        "suffix": (
            ", modern architecture, glass facade, concrete texture, minimalist, "
            "clean lines, architectural visualization, "
            "golden hour, reflection, dramatic light, "
            "trending on archdaily, behance, unreal engine 5 render, "
            "highly detailed, 8k"
        ),
        "negative": (
            "low quality, blurry, distorted architecture, mirror error, "
            "floating mass, readable text, watermark, logo, "
            "messy facade, bad scale"
        ),
    },
    "char": {
        "name": "人物角色",
        "suffix": (
            ", high-end fashion lookbook, character style sheet, "
            "professional studio lighting, photorealistic, "
            "four views of the same character in one row, "
            "close-up portrait, full-length front, three-quarter side, back view, "
            "plain light warm-grey seamless studio background, "
            "ultra detailed, 8k"
        ),
        "negative": (
            "low quality, blurry, extra fingers, bad anatomy, deformed face, "
            "asymmetry, wrong proportions, text, watermark, logo, border, "
            "cartoon, painting style, duplicate face"
        ),
    },
    "wabi": {
        "name": "日式侘寂",
        "suffix": (
            ", Japanese wabi-sabi architecture, exposed concrete, wood lattice, "
            "zen garden, soft natural light, minimalist, "
            "muted earth tones, imperfect beauty, "
            "architectural photography, cinematic lighting, "
            "highly detailed, 8k"
        ),
        "negative": (
            "low quality, blurry, bright colors, cluttered, "
            "readable text, watermark, logo, modern technology"
        ),
    },
    "cyber": {
        "name": "赛博朋克",
        "suffix": (
            ", cyberpunk cityscape, neon lights, glass curtain wall, "
            "futuristic, night scene, rain reflections, "
            "holographic advertisements, flying vehicles, "
            "cinematic lighting, volumetric fog, "
            "highly detailed, 8k, trending on artstation"
        ),
        "negative": (
            "low quality, blurry, daylight, rural, "
            "readable text, watermark, logo"
        ),
    },
    "ink": {
        "name": "水墨风格",
        "suffix": (
            ", Chinese ink wash painting style, sumi-e, "
            "monochrome black and white, brush strokes, "
            "traditional Chinese landscape, mountains and rivers, "
            "minimalist composition, rice paper texture, "
            "artistic, elegant, 8k"
        ),
        "negative": (
            "low quality, blurry, colorful, photographic, "
            "readable text, watermark, logo, modern elements"
        ),
    },
}

# === 尺寸预设 ===
SIZES = {
    "1024x1024": "1:1 方图",
    "1024x1792": "9:16 竖图",
    "1792x1024": "16:9 横图",
    "1344x768": "16:9 宽屏",
    "768x1344": "9:16 竖屏",
}


def get_api_key():
    key = os.environ.get("AGNES_API_KEY")
    if key:
        return key
    env_file = Path.home() / ".hermes" / ".env"
    if env_file.exists():
        for line in env_file.read_text().splitlines():
            if line.startswith("AGNES_API_KEY="):
                return line.split("=", 1)[1].strip()
    print("❌ 未找到 AGNES_API_KEY")
    sys.exit(1)


def curl_post(url, data, timeout=120):
    """用curl发POST请求（避免urllib IPv6问题）"""
    key = get_api_key()
    cmd = ["curl", "-s", "--max-time", str(timeout), "-X", "POST", url,
           "-H", f"Authorization: Bearer {key}",
           "-H", "Content-Type: application/json",
           "-d", json.dumps(data)]
    env = dict(os.environ)
    for k in ["http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"]:
        env.pop(k, None)
    r = subprocess.run(cmd, capture_output=True, timeout=timeout+10, env=env)
    text = r.stdout.decode("utf-8", errors="ignore")
    if not text.strip():
        return None
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return {"raw": text}


def curl_download(url, output_path, timeout=60):
    """用curl下载文件"""
    cmd = ["curl", "-s", "--max-time", str(timeout), "-o", str(output_path), url]
    env = dict(os.environ)
    for k in ["http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"]:
        env.pop(k, None)
    r = subprocess.run(cmd, capture_output=True, timeout=timeout+5, env=env)
    return r.returncode == 0 and Path(output_path).exists()


def generate_image(prompt, style="default", size="1024x1024", output=None, push=False):
    """生成图片"""
    style_cfg = STYLES.get(style, STYLES["default"])

    # 拼接完整 prompt
    full_prompt = prompt + style_cfg["suffix"]

    print(f"🎨 Agnes 图片生成")
    print(f"   模型: {MODEL}")
    print(f"   风格: {style_cfg['name']}")
    print(f"   尺寸: {size}")
    print(f"   提示词: {prompt}")
    print()

    # 调用 API
    data = {
        "model": MODEL,
        "prompt": full_prompt,
        "n": 1,
        "size": size,
    }

    print("⏳ 生成中（约30-60秒）...")
    start = time.time()
    result = curl_post(f"{API_BASE}/images/generations", data, timeout=120)

    if not result:
        # 尝试备用模型
        print(f"⚠️ {MODEL} 超时，尝试 {FALLBACK_MODEL}...")
        data["model"] = FALLBACK_MODEL
        result = curl_post(f"{API_BASE}/images/generations", data, timeout=120)

    if not result:
        print("❌ 生成失败（API 无响应）")
        return None

    # 解析响应
    image_url = None
    if "data" in result and result["data"]:
        image_url = result["data"][0].get("url") or result["data"][0].get("b64_json")
    elif "url" in result:
        image_url = result["url"]
    elif "error" in result:
        print(f"❌ API 错误: {result['error']}")
        return None

    if not image_url:
        print(f"❌ 未找到图片 URL")
        print(f"   响应: {json.dumps(result, indent=2)[:500]}")
        return None

    # 如果是 base64，保存为文件
    if image_url.startswith("data:"):
        import base64
        header, b64data = image_url.split(",", 1)
        ext = "png" if "png" in header else "jpg"
        if not output:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output = str(OUTPUT_DIR / f"agnes_{style}_{timestamp}.{ext}")
        Path(output).write_bytes(base64.b64decode(b64data))
    else:
        # 下载图片
        if not output:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output = str(OUTPUT_DIR / f"agnes_{style}_{timestamp}.png")
        if not curl_download(image_url, output):
            print("❌ 下载失败")
            return None

    elapsed = time.time() - start
    size_kb = Path(output).stat().st_size / 1024
    print(f"\n✅ 生成完成！({elapsed:.0f}秒)")
    print(f"📁 已保存: {output} ({size_kb:.0f}KB)")
    print(f"   风格: {style_cfg['name']}")

    # 推飞书
    if push:
        send_feishu(output, prompt, style_cfg["name"])

    return output


def send_feishu(image_path, prompt, style_name):
    """推送到飞书"""
    msg = f"""🎨 Agnes 图片生成完成

风格: {style_name}
提示词: {prompt}
文件: {image_path}

MEDIA:{image_path}"""

    payload = json.dumps({"msg_type": "text", "content": {"text": msg}}).encode()
    cmd = ["curl", "-s", "--max-time", "10", "-X", "POST", FEISHU_WEBHOOK,
           "-H", "Content-Type: application/json", "-d", payload.decode()]
    env = dict(os.environ)
    for k in ["http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"]:
        env.pop(k, None)
    subprocess.run(cmd, capture_output=True, timeout=12, env=env)
    print("📤 已推送飞书")


def list_cache():
    """列出本地缓存"""
    images = sorted(OUTPUT_DIR.glob("agnes_*.png"), reverse=True)
    images += sorted(OUTPUT_DIR.glob("agnes_*.jpg"), reverse=True)
    if not images:
        print("📭 暂无缓存图片")
        return
    print(f"📁 图片缓存 ({len(images)}张):")
    for img in images[:20]:
        size_kb = img.stat().st_size / 1024
        mtime = datetime.fromtimestamp(img.stat().st_mtime).strftime("%m-%d %H:%M")
        print(f"   {mtime} {size_kb:.0f}KB {img.name}")


def list_styles():
    """列出所有风格"""
    print("🎨 可用风格:")
    for key, cfg in STYLES.items():
        print(f"   --style {key:10s} {cfg['name']}")
    print()
    print("📐 可用尺寸:")
    for size, desc in SIZES.items():
        print(f"   --size {size:12s} {desc}")


def main():
    if len(sys.argv) < 2:
        print("用法:")
        print('  python3 agnes_image.py "提示词" [--style 风格] [--size 尺寸] [--output 路径] [--push]')
        print("  python3 agnes_image.py list      # 列出缓存")
        print("  python3 agnes_image.py styles     # 列出风格")
        return

    if sys.argv[1] == "list":
        list_cache()
        return

    if sys.argv[1] == "styles":
        list_styles()
        return

    # 解析参数
    prompt = sys.argv[1]
    style = "default"
    size = "1024x1024"
    output = None
    push = False

    args = sys.argv[2:]
    i = 0
    while i < len(args):
        if args[i] == "--style" and i + 1 < len(args):
            style = args[i + 1]; i += 2
        elif args[i] == "--size" and i + 1 < len(args):
            size = args[i + 1]; i += 2
        elif args[i] == "--output" and i + 1 < len(args):
            output = args[i + 1]; i += 2
        elif args[i] == "--push":
            push = True; i += 1
        else:
            i += 1

    generate_image(prompt, style, size, output, push)


if __name__ == "__main__":
    main()