#!/usr/bin/env python3
"""Gemini の画像 API でマンガの1コマを1枚だけ生成する。

使い方:
  python3 manga/gemini_panel.py --prompt P.txt --image canvas.png --image hakobu.png \
      --out-dir OUT --model gemini-3.1-flash-image --name 第61話_P1

- 鍵は /Volumes/KIOXIA/gemini_api_key.txt から読むだけ (表示・保存しない)
- 縦横比 4:5 / 大きさ 1K を generationConfig.imageConfig で指定
- 生成は1回だけ。失敗してもやり直さない
- 出力: {name}_api_{model}_{日時}.{png|jpg|webp} (返ってきた形式に合わせる) と同じ名前の .json
  (使用量など。画像データは含めない)
"""
import argparse
import json
import mimetypes
import sys
import time
from datetime import datetime
from pathlib import Path

from google import genai
from google.genai import types

KEY_PATH = Path("/Volumes/KIOXIA/gemini_api_key.txt")
EXT_BY_MIME = {"image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp"}


def load_key():
    if not Path("/Volumes/KIOXIA").is_dir():
        sys.exit("ABORT: KIOXIA がつながっていない")
    if not KEY_PATH.is_file():
        sys.exit(f"ABORT: 鍵のファイルがない: {KEY_PATH}")
    key = KEY_PATH.read_text().strip()
    if not key:
        sys.exit("ABORT: 鍵のファイルが空")
    return key


def strip_inline_data(obj):
    """返事の dict から画像のバイト列を取り除く。"""
    if isinstance(obj, dict):
        out = {}
        for k, v in obj.items():
            if k in ("inline_data", "inlineData") and isinstance(v, dict):
                out[k] = {kk: vv for kk, vv in v.items() if kk != "data"}
                out[k]["data"] = "<omitted>"
            elif k in ("thought_signature", "thoughtSignature"):
                out[k] = "<omitted>"
            else:
                out[k] = strip_inline_data(v)
        return out
    if isinstance(obj, list):
        return [strip_inline_data(v) for v in obj]
    return obj


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt", required=True, type=Path)
    ap.add_argument("--image", action="append", default=[], type=Path,
                    help="添付する画像 (この順で渡す。複数可)")
    ap.add_argument("--out-dir", required=True, type=Path)
    ap.add_argument("--model", required=True)
    ap.add_argument("--name", default="panel", help="出力ファイル名の頭")
    ap.add_argument("--aspect-ratio", default="4:5")
    ap.add_argument("--image-size", default="1K")
    args = ap.parse_args()

    prompt = args.prompt.read_text(encoding="utf-8")
    if not prompt.strip():
        sys.exit(f"ABORT: プロンプトが空: {args.prompt}")
    for p in args.image:
        if not p.is_file():
            sys.exit(f"ABORT: 画像がない: {p}")

    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    stem = f"{args.name}_api_{args.model}_{stamp}"
    args.out_dir.mkdir(parents=True, exist_ok=True)
    json_path = args.out_dir / f"{stem}.json"
    if any(args.out_dir.glob(f"{stem}.*")):
        sys.exit(f"ABORT: 同名ファイルがある: {stem}")

    contents = [prompt]
    for p in args.image:
        mime = mimetypes.guess_type(p.name)[0] or "image/png"
        contents.append(types.Part.from_bytes(data=p.read_bytes(), mime_type=mime))

    config = types.GenerateContentConfig(
        response_modalities=["TEXT", "IMAGE"],
        image_config=types.ImageConfig(
            aspect_ratio=args.aspect_ratio,
            image_size=args.image_size,
        ),
    )

    client = genai.Client(api_key=load_key())
    record = {
        "requested_model": args.model,
        "prompt_file": str(args.prompt),
        "images": [str(p) for p in args.image],
        "aspect_ratio": args.aspect_ratio,
        "image_size": args.image_size,
        "started_at": datetime.now().isoformat(timespec="seconds"),
    }
    t0 = time.time()
    try:
        resp = client.models.generate_content(
            model=args.model, contents=contents, config=config)
    except Exception as e:
        record["elapsed_sec"] = round(time.time() - t0, 2)
        record["error"] = f"{type(e).__name__}: {e}"
        json_path.write_text(json.dumps(record, ensure_ascii=False, indent=2),
                             encoding="utf-8")
        print(f"ERROR: {record['error']}")
        print(f"json: {json_path}")
        sys.exit(1)
    record["elapsed_sec"] = round(time.time() - t0, 2)

    image_saved = False
    texts = []
    for cand in resp.candidates or []:
        for part in (cand.content.parts if cand.content else None) or []:
            if part.inline_data and part.inline_data.data and not image_saved:
                mime = part.inline_data.mime_type or "image/png"
                img_path = args.out_dir / f"{stem}{EXT_BY_MIME.get(mime, '.bin')}"
                img_path.write_bytes(part.inline_data.data)
                record["image_mime_type"] = mime
                image_saved = True
            elif part.text and not part.thought:
                texts.append(part.text)

    record["model_version"] = resp.model_version
    record["response_id"] = resp.response_id
    record["usage_metadata"] = (resp.usage_metadata.model_dump(mode="json", exclude_none=True)
                                if resp.usage_metadata else None)
    record["finish_reasons"] = [str(c.finish_reason) for c in resp.candidates or []]
    record["texts"] = texts
    record["image_file"] = str(img_path) if image_saved else None
    record["response"] = strip_inline_data(resp.model_dump(mode="json", exclude_none=True))
    json_path.write_text(json.dumps(record, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"elapsed: {record['elapsed_sec']} s / model_version: {resp.model_version}")
    print(f"usage: {json.dumps(record['usage_metadata'], ensure_ascii=False)}")
    print(f"image: {img_path if image_saved else '(画像なし)'}")
    print(f"json: {json_path}")
    if not image_saved:
        sys.exit(1)


if __name__ == "__main__":
    main()