#!/usr/bin/env python3
"""MAOMOMO article image generation CLI.

This is a lightweight OpenAI-compatible fallback for article assets. Prefer the
agent's built-in image tool when it is available; use this script when the user
explicitly wants API/CLI generation or the built-in backend is unavailable.
"""

from __future__ import annotations

import argparse
import base64
import json
import os
from pathlib import Path
import sys
import time
from typing import Any, Dict, Iterable, List, Optional
import urllib.error
import urllib.request


DEFAULT_BASE_URL = "https://api.openai.com/v1"
DEFAULT_MODEL = "gpt-image-2"
DEFAULT_SIZE = "1536x1024"
DEFAULT_QUALITY = "medium"
DEFAULT_STYLE = "warm-fintech-guide"


STYLE_PRESETS: Dict[str, str] = {
    "warm-fintech-guide": (
        "MAOMOMO 暖橙金融教程风：明亮白底，暖橙强调色，干净卡片布局，友好的白橘猫作为向导，"
        "适合港卡、银行活动、返现攻略和教程总览。"
    ),
    "clean-editorial": (
        "清爽编辑部风：大标题、留白充足、少量暖橙和黑灰文字，像一张信息密度适中的中文攻略头图。"
    ),
    "data-card-dashboard": (
        "数据卡片仪表盘风：指标卡、时间线、计算公式和对比表清晰分区，适合返现、费用、门槛和路径对比。"
    ),
    "handdrawn-note": (
        "手绘便签风：白底纸感、手绘箭头、便利贴、重点圈注和轻量猫咪贴纸，适合避坑经验和保姆式步骤。"
    ),
    "xiaohongshu-vertical": (
        "小红书竖版攻略风：9:16 手机端可读，大字少字，强标题钩子，4 图系列一致视觉，带 MAOMOMO 标识。"
    ),
    "clean-professional": (
        "清爽专业风：浅色背景、蓝绿或暖橙强调、结构化信息卡片和清晰层级，适合正式银行规则、开户教程和综合攻略。"
    ),
    "creative-magazine": (
        "创意杂志风：大标题、强留白、编辑部排版和轻视觉冲击，适合传播型封面、观点总结和活动盘点。"
    ),
    "retro-flat-illustration": (
        "复古扁平插画风：低饱和暖色、扁平金融小物件、轻复古海报感，适合轻松羊毛攻略和经验分享。"
    ),
    "e-ink-editorial": (
        "电子墨水杂志风：纸感背景、黑白灰为主、少量强调色、强标题和元信息条，适合深度解释和观点型长文。"
    ),
    "scientific-defense": (
        "科研答辩风：严谨浅色版式、证据图、流程框和来源标注，适合规则拆解、政策解读和多来源核对。"
    ),
    "mckinsey-brief": (
        "麦肯锡简报风：结论先行、矩阵、2x2、瀑布图和高对比商业配色，适合方案对比、路径选择和决策建议。"
    ),
}


def die(message: str, code: int = 1) -> None:
    print(f"错误：{message}", file=sys.stderr)
    raise SystemExit(code)


def warn(message: str) -> None:
    print(f"警告：{message}", file=sys.stderr)


def read_text(path: Path) -> str:
    return path.read_text(encoding="utf-8")


def write_text(path: Path, text: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="\n") as handle:
        handle.write(text)


def read_json(path: Path) -> Any:
    return json.loads(read_text(path))


def load_codex_config_base_url() -> Optional[str]:
    config_path = Path.home() / ".codex" / "config.toml"
    if not config_path.exists():
        return None
    try:
        import tomllib
    except ModuleNotFoundError:
        return None
    try:
        data = tomllib.loads(read_text(config_path))
    except Exception:
        return None

    provider_name = data.get("model_provider")
    providers = data.get("model_providers")
    if isinstance(provider_name, str) and isinstance(providers, dict):
        provider = providers.get(provider_name)
        if isinstance(provider, dict) and isinstance(provider.get("base_url"), str):
            return provider["base_url"].strip()

    if isinstance(data.get("base_url"), str):
        return data["base_url"].strip()
    if isinstance(providers, dict):
        for provider in providers.values():
            if isinstance(provider, dict) and isinstance(provider.get("base_url"), str):
                return provider["base_url"].strip()
    return None


def find_secret(value: Any) -> Optional[str]:
    if isinstance(value, dict):
        for key in ("MAOMOMO_IMAGE_API_KEY", "OPENAI_API_KEY", "api_key", "openai_api_key", "token"):
            raw = value.get(key)
            if isinstance(raw, str) and raw.strip():
                return raw.strip()
        for child in value.values():
            found = find_secret(child)
            if found:
                return found
    if isinstance(value, list):
        for child in value:
            found = find_secret(child)
            if found:
                return found
    return None


def load_codex_auth_api_key() -> Optional[str]:
    auth_path = Path.home() / ".codex" / "auth.json"
    if not auth_path.exists():
        return None
    try:
        return find_secret(read_json(auth_path))
    except Exception:
        return None


def api_key() -> str:
    value = os.getenv("MAOMOMO_IMAGE_API_KEY") or os.getenv("OPENAI_API_KEY") or load_codex_auth_api_key()
    if not value:
        die(
            "未找到 API Key。请设置 MAOMOMO_IMAGE_API_KEY 或 OPENAI_API_KEY，"
            "或确保 ~/.codex/auth.json 中存在可用 key。"
        )
    return value


def base_url() -> str:
    return (
        os.getenv("MAOMOMO_IMAGE_BASE_URL")
        or os.getenv("OPENAI_BASE_URL")
        or load_codex_config_base_url()
        or DEFAULT_BASE_URL
    ).rstrip("/")


def image_model() -> str:
    return os.getenv("MAOMOMO_IMAGE_MODEL") or os.getenv("CODEX_PPT_IMAGE_MODEL") or DEFAULT_MODEL


def style_prompt(style: str) -> str:
    if style in STYLE_PRESETS:
        return STYLE_PRESETS[style]
    return style


def build_prompt(
    *,
    title: str,
    image_type: str,
    core_text: str,
    style: str,
    aspect_ratio: str,
    extra: str = "",
) -> str:
    parts = [
        "生成一张 MAOMOMO 中文金融实操文章配图。",
        f"图片类型：{image_type}",
        f"标题 / 主题：{title}",
        f"核心文案：{core_text}",
        f"画幅：{aspect_ratio}",
        f"视觉风格：{style_prompt(style)}",
        "硬性要求：包含清晰可见的 MAOMOMO 标识；中文文字清楚可读；信息层级明确；不要伪造真实 App 截图；不要使用未经提供的银行、券商、支付机构或卡组织 Logo；不要出现真实个人信息；猫咪不能承载第三方品牌 Logo。",
    ]
    if extra.strip():
        parts.append(f"补充要求：{extra.strip()}")
    return "\n".join(parts)


def request_image(prompt: str, *, size: str, quality: str, output_format: str) -> bytes:
    payload = {
        "model": image_model(),
        "prompt": prompt,
        "size": size,
        "quality": quality,
        "n": 1,
        "response_format": "b64_json",
    }
    if output_format:
        payload["output_format"] = output_format

    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        base_url() + "/images/generations",
        data=body,
        headers={
            "Authorization": f"Bearer {api_key()}",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=180) as resp:
            data = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")
        die(f"图片接口返回 HTTP {exc.code}: {detail}")
    except Exception as exc:
        die(f"图片接口请求失败：{exc}")

    items = data.get("data")
    if not isinstance(items, list) or not items:
        die("图片接口响应缺少 data。")
    first = items[0]
    if not isinstance(first, dict):
        die("图片接口响应格式不正确。")
    if isinstance(first.get("b64_json"), str):
        return base64.b64decode(first["b64_json"])
    if isinstance(first.get("url"), str):
        with urllib.request.urlopen(first["url"], timeout=180) as resp:
            return resp.read()
    die("图片接口响应中没有 b64_json 或 url。")
    return b""


def normalize_output(path: str, *, base_dir: Optional[Path]) -> Path:
    out = Path(path)
    if not out.is_absolute() and base_dir is not None:
        out = base_dir / out
    return out.resolve()


def generate_one(args: argparse.Namespace) -> Dict[str, str]:
    if args.prompt_file:
        prompt = read_text(Path(args.prompt_file)).strip()
    elif args.prompt:
        prompt = args.prompt.strip()
    else:
        prompt = build_prompt(
            title=args.title,
            image_type=args.image_type,
            core_text=args.core_text,
            style=args.style,
            aspect_ratio=args.aspect_ratio,
            extra=args.extra,
        )

    out = normalize_output(args.out, base_dir=None)
    if args.dry_run:
        write_text(out.with_suffix(".prompt.txt"), prompt + "\n")
        return {"out": str(out), "prompt": str(out.with_suffix(".prompt.txt")), "status": "dry-run"}

    image_bytes = request_image(prompt, size=args.size, quality=args.quality, output_format=args.output_format)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_bytes(image_bytes)
    return {"out": str(out), "status": "generated"}


def manifest_items(path: Path) -> List[Dict[str, Any]]:
    data = read_json(path)
    if isinstance(data, dict):
        items = data.get("images")
    else:
        items = data
    if not isinstance(items, list):
        die("manifest 必须是图片数组，或包含 images 数组的对象。")
    normalized: List[Dict[str, Any]] = []
    for index, item in enumerate(items, start=1):
        if not isinstance(item, dict):
            die(f"manifest 第 {index} 项不是对象。")
        normalized.append(item)
    return normalized


def batch(args: argparse.Namespace) -> List[Dict[str, str]]:
    manifest_path = Path(args.manifest).resolve()
    base_dir = Path(args.base_dir).resolve() if args.base_dir else manifest_path.parent
    results: List[Dict[str, str]] = []

    for index, item in enumerate(manifest_items(manifest_path), start=1):
        filename = item.get("file_name") or item.get("filename") or item.get("out")
        if not isinstance(filename, str) or not filename.strip():
            die(f"manifest 第 {index} 项缺少 file_name。")
        style = str(item.get("style") or args.style or DEFAULT_STYLE)
        aspect_ratio = str(item.get("aspect_ratio") or args.aspect_ratio)
        prompt = item.get("prompt")
        if not isinstance(prompt, str) or not prompt.strip():
            prompt = build_prompt(
                title=str(item.get("title") or item.get("alt_text") or filename),
                image_type=str(item.get("type") or "文章配图"),
                core_text=str(item.get("core_text") or item.get("caption") or item.get("alt_text") or ""),
                style=style,
                aspect_ratio=aspect_ratio,
                extra=str(item.get("extra") or item.get("requirements") or ""),
            )

        out = normalize_output(filename, base_dir=base_dir)
        if args.dry_run:
            write_text(out.with_suffix(".prompt.txt"), prompt + "\n")
            results.append({"out": str(out), "prompt": str(out.with_suffix(".prompt.txt")), "status": "dry-run"})
            continue

        image_bytes = request_image(
            prompt,
            size=str(item.get("size") or args.size),
            quality=str(item.get("quality") or args.quality),
            output_format=str(item.get("output_format") or args.output_format),
        )
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_bytes(image_bytes)
        results.append({"out": str(out), "status": "generated"})
        if args.sleep > 0:
            time.sleep(args.sleep)
    return results


def print_styles(_: argparse.Namespace) -> int:
    print(json.dumps(STYLE_PRESETS, ensure_ascii=False, indent=2))
    return 0


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Generate MAOMOMO article images with an OpenAI-compatible API.")
    subparsers = parser.add_subparsers(dest="command", required=True)

    common = argparse.ArgumentParser(add_help=False)
    common.add_argument("--style", default=DEFAULT_STYLE, help="Style preset name or custom style text.")
    common.add_argument("--size", default=DEFAULT_SIZE)
    common.add_argument("--quality", default=DEFAULT_QUALITY)
    common.add_argument("--output-format", default="png")
    common.add_argument("--aspect-ratio", default="16:9")
    common.add_argument("--dry-run", action="store_true", help="Only write prompt files; do not call the image API.")

    generate = subparsers.add_parser("generate", parents=[common], help="Generate one image.")
    generate.add_argument("--out", required=True)
    generate.add_argument("--prompt")
    generate.add_argument("--prompt-file")
    generate.add_argument("--title", default="MAOMOMO 文章配图")
    generate.add_argument("--image-type", default="文章配图")
    generate.add_argument("--core-text", default="")
    generate.add_argument("--extra", default="")

    batch_parser = subparsers.add_parser("batch", parents=[common], help="Generate images from a JSON manifest.")
    batch_parser.add_argument("--manifest", required=True)
    batch_parser.add_argument("--base-dir", help="Base directory for relative file_name paths. Defaults to manifest dir.")
    batch_parser.add_argument("--sleep", type=float, default=0.0, help="Seconds to sleep between API calls.")

    subparsers.add_parser("styles", help="Print bundled style presets.")
    return parser


def main(argv: Optional[Iterable[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(list(argv) if argv is not None else None)
    if args.command == "styles":
        return print_styles(args)
    if args.command == "generate":
        result = generate_one(args)
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return 0
    if args.command == "batch":
        results = batch(args)
        print(json.dumps(results, ensure_ascii=False, indent=2))
        return 0
    parser.print_help()
    return 1


if __name__ == "__main__":
    raise SystemExit(main())