"""Generate or edit images through the Codex CLI (ChatGPT subscription quota).

Codex saves images under ~/.codex/generated_images/; this script collects the
PNGs created during the run and copies them to the output folder.

Examples:
    python scripts/openai-image/codex_image.py "一隻戴眼鏡的橘貓老師，扁平插畫風" --aspect landscape
    python scripts/openai-image/codex_image.py "把背景改成夜晚" --edit input.png
"""

import argparse
import glob
import os
import re
import shutil
import subprocess
import sys
import tempfile
import time
from datetime import datetime
from pathlib import Path

GENERATED_DIR = Path.home() / ".codex" / "generated_images"
ASPECTS = {
    "landscape": "landscape (wide, about 3:2)",
    "portrait": "portrait (tall, about 2:3)",
    "square": "square (1:1)",
}


def find_codex():
    override = os.environ.get("CODEX_EXE")
    if override and Path(override).exists():
        return override
    on_path = shutil.which("codex")
    if on_path:
        return on_path
    # Codex desktop app (Microsoft Store) bundles the CLI; its folder name changes per version.
    try:
        loc = subprocess.run(
            ["powershell", "-NoProfile", "-Command",
             "(Get-AppxPackage *OpenAI.Codex* | Sort-Object Version -Descending | Select-Object -First 1).InstallLocation"],
            capture_output=True, text=True, timeout=30,
        ).stdout.strip()
    except (OSError, subprocess.TimeoutExpired):
        loc = ""
    if loc:
        exe = Path(loc) / "app" / "resources" / "codex.exe"
        if exe.exists():
            return str(exe)
    matches = sorted(glob.glob(r"C:\Program Files\WindowsApps\OpenAI.Codex_*\app\resources\codex.exe"))
    if matches:
        return matches[-1]
    sys.exit("找不到 Codex CLI。請確認已安裝 Codex 桌面版或 codex CLI，或設定 CODEX_EXE 環境變數。")


def slugify(text, max_len=30):
    slug = re.sub(r"[^\w\u4e00-\u9fff-]+", "-", text).strip("-")
    return slug[:max_len] or "image"


def parse_args():
    parser = argparse.ArgumentParser(description="透過 Codex CLI 生圖 / 修圖")
    parser.add_argument("prompt", nargs="?", help="圖片描述（中英文皆可）")
    parser.add_argument("--prompt-file", type=Path,
                        help="從 UTF-8 文字檔讀取描述（長描述或含引號時使用）")
    parser.add_argument("--aspect", choices=ASPECTS, help="landscape / portrait / square")
    parser.add_argument("--edit", nargs="+", type=Path, metavar="IMAGE",
                        help="附上既有圖片作為修改底圖或參考圖")
    parser.add_argument("-n", type=int, default=1, help="張數")
    parser.add_argument("--out-dir", type=Path, default=Path.cwd() / "output" / "images")
    parser.add_argument("--name", help="輸出檔名前綴（預設取自 prompt）")
    parser.add_argument("--timeout", type=int, default=600, help="秒")
    args = parser.parse_args()
    if args.prompt_file:
        args.prompt = args.prompt_file.read_text(encoding="utf-8").strip()
    if not args.prompt:
        parser.error("請提供 prompt 或 --prompt-file")
    return args


def build_instruction(args):
    lines = [
        f"Use your image generation tool to create {args.n} image(s).",
        "Do not write code, do not edit any files, and do not run shell commands.",
    ]
    if args.edit:
        lines.append("Use the attached image(s) as the base/reference and apply the requested changes.")
    if args.aspect:
        lines.append(f"Aspect ratio: {ASPECTS[args.aspect]}.")
    lines.append("Image request (follow it faithfully; render any quoted text exactly):")
    lines.append(args.prompt)
    lines.append("When done, reply only with DONE.")
    return "\n".join(lines)


def main():
    args = parse_args()
    codex = find_codex()
    cmd = [codex, "exec", "--skip-git-repo-check", "-s", "read-only",
           "-C", tempfile.gettempdir()]
    for img in args.edit or []:
        cmd += ["-i", str(img.resolve())]
    cmd.append("-")

    started = time.time() - 1
    try:
        proc = subprocess.run(cmd, input=build_instruction(args), capture_output=True,
                              text=True, encoding="utf-8", errors="replace",
                              timeout=args.timeout)
    except subprocess.TimeoutExpired:
        sys.exit(f"Codex 逾時（{args.timeout} 秒）。")

    new_images = sorted(
        (p for p in GENERATED_DIR.rglob("*.png") if p.stat().st_mtime >= started),
        key=lambda p: p.stat().st_mtime,
    )
    if not new_images:
        sys.stderr.write(proc.stdout[-2000:] + proc.stderr[-2000:])
        sys.exit("Codex 沒有產生圖片（可能是額度用完、未登入或請求被拒）。")

    args.out_dir.mkdir(parents=True, exist_ok=True)
    stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    prefix = args.name or slugify(args.prompt.splitlines()[0])
    for i, src in enumerate(new_images, start=1):
        suffix = f"-{i}" if len(new_images) > 1 else ""
        dst = args.out_dir / f"{stamp}-{prefix}{suffix}.png"
        shutil.copy2(src, dst)
        print(dst)


if __name__ == "__main__":
    main()