"""Generate or edit images with the OpenAI Images API.

API key lookup order: OPENAI_API_KEY env var, then a `.env` file at the
repository root, next to this script, or in the current directory
(keep it git-ignored). Never pass the key on the command line.

Examples:
    python scripts/openai-image/generate_image.py "一隻戴眼鏡的橘貓在教室講課，扁平插畫風"
    python scripts/openai-image/generate_image.py "改成夜晚場景" --edit input.png
    python scripts/openai-image/generate_image.py "Logo" --size 1024x1024 --background transparent -n 2
"""

import argparse
import base64
import os
import re
import sys
from datetime import datetime
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_OUT_DIR = Path.cwd() / "output" / "images"
DEFAULT_MODEL = os.environ.get("OPENAI_IMAGE_MODEL", "gpt-image-2.5-flare")


def load_api_key():
    key = os.environ.get("OPENAI_API_KEY")
    if key:
        return key
    script_dir = Path(__file__).resolve().parent
    for env_file in (REPO_ROOT / ".env", script_dir / ".env", Path.cwd() / ".env"):
        if not env_file.exists():
            continue
        for line in env_file.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if line.startswith("OPENAI_API_KEY="):
                return line.split("=", 1)[1].strip().strip('"').strip("'")
    sys.exit(
        "找不到 OPENAI_API_KEY。請設定環境變數，或在腳本旁邊的 .env 寫入：\n"
        "OPENAI_API_KEY=sk-..."
    )


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


def parse_args():
    parser = argparse.ArgumentParser(description="OpenAI 圖片生成 / 編輯")
    parser.add_argument("prompt", help="圖片描述（中英文皆可）")
    parser.add_argument("--edit", nargs="+", type=Path, metavar="IMAGE",
                        help="以既有圖片為底進行編輯（可多張參考圖）")
    parser.add_argument("--mask", type=Path, help="編輯用遮罩 PNG（透明處為要修改的區域）")
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--size", default="auto",
                        help="1024x1024 / 1536x1024（橫）/ 1024x1536（直）/ auto")
    parser.add_argument("--quality", default="auto", help="low / medium / high / auto")
    parser.add_argument("--background", default="auto", help="transparent / opaque / auto")
    parser.add_argument("--format", default="png", choices=["png", "jpeg", "webp"])
    parser.add_argument("-n", type=int, default=1, help="張數")
    parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT_DIR)
    parser.add_argument("--name", help="輸出檔名前綴（預設取自 prompt）")
    return parser.parse_args()


def main():
    args = parse_args()
    from openai import OpenAI

    client = OpenAI(api_key=load_api_key())
    common = dict(model=args.model, prompt=args.prompt, n=args.n, size=args.size,
                  quality=args.quality)

    if args.edit:
        files = [open(p, "rb") for p in args.edit]
        try:
            extra = {"mask": open(args.mask, "rb")} if args.mask else {}
            result = client.images.edit(image=files if len(files) > 1 else files[0],
                                        background=args.background, **extra, **common)
        finally:
            for f in files:
                f.close()
    else:
        result = client.images.generate(background=args.background,
                                        output_format=args.format, **common)

    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)
    ext = "png" if args.edit else args.format

    for i, item in enumerate(result.data, start=1):
        suffix = f"-{i}" if len(result.data) > 1 else ""
        path = args.out_dir / f"{stamp}-{prefix}{suffix}.{ext}"
        path.write_bytes(base64.b64decode(item.b64_json))
        print(path)

    usage = getattr(result, "usage", None)
    if usage:
        print(f"tokens: input={usage.input_tokens} output={usage.output_tokens}", file=sys.stderr)


if __name__ == "__main__":
    main()