scripts/batch_generate_images.py

#!/usr/bin/env python3
"""
일괄 이미지 생성 스크립트
Gemini gemini-2.5-flash-image (나노바나나) 모델로 스팟 이미지 일괄 생성

사용법:
    # 모든 카테고리 폴백 이미지 생성
    python scripts/batch_generate_images.py --fallback

    # 특정 카테고리 스팟 이미지 생성 (최대 5개)
    python scripts/batch_generate_images.py --category beach --limit 5

    # generated_image_url이 없는 모든 스팟 (최대 10개)
    python scripts/batch_generate_images.py --limit 10

    # 특정 스팟 ID로 생성
    python scripts/batch_generate_images.py --spot-ids spot1 spot2 spot3

    # 드라이런 (실제 생성 없이 대상 목록만 확인)
    python scripts/batch_generate_images.py --dry-run --limit 20

환경변수:
    GEMINI_API_KEY 또는 GOOGLE_API_KEY: Gemini API 키
    SUPABASE_URL: Supabase 프로젝트 URL
    SUPABASE_SERVICE_ROLE_KEY: Supabase service role 키
"""
import argparse
import asyncio
import base64
import logging
import os
import sys
import time

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from dotenv import load_dotenv
load_dotenv()

from db import get_supabase

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)

# Import prompts from the router
from routers.image_gen import CATEGORY_PROMPTS, DEFAULT_PROMPT, NANO_BANANA_MODEL, STORAGE_BUCKET


def get_genai_client():
    """Get Google GenAI client"""
    from google import genai

    api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY")
    if not api_key:
        raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY environment variable is required")

    return genai.Client(api_key=api_key)


def generate_image(client, prompt: str) -> tuple[bytes | None, str]:
    """Generate image with Gemini and return (image_bytes, mime_type)"""
    from google.genai import types

    response = client.models.generate_content(
        model=NANO_BANANA_MODEL,
        contents=[prompt],
        config=types.GenerateContentConfig(
            response_modalities=["IMAGE", "TEXT"],
        ),
    )

    if response.candidates and response.candidates[0].content:
        for part in response.candidates[0].content.parts:
            if hasattr(part, "inline_data") and part.inline_data:
                data = part.inline_data.data
                mime = part.inline_data.mime_type or "image/webp"
                image_bytes = data if isinstance(data, bytes) else base64.b64decode(data)
                return image_bytes, mime

    return None, ""


def upload_to_storage(supabase, path: str, image_bytes: bytes, mime_type: str) -> str:
    """Upload image to Supabase Storage and return public URL"""
    try:
        supabase.storage.from_(STORAGE_BUCKET).remove([path])
    except Exception:
        pass

    supabase.storage.from_(STORAGE_BUCKET).upload(
        path=path,
        file=image_bytes,
        file_options={"content-type": mime_type, "upsert": "true"},
    )

    return supabase.storage.from_(STORAGE_BUCKET).get_public_url(path)


def generate_fallback_images(dry_run: bool = False):
    """Generate category fallback images"""
    logger.info(f"Generating fallback images for {len(CATEGORY_PROMPTS)} categories...")

    if dry_run:
        for category in CATEGORY_PROMPTS:
            logger.info(f"  [DRY-RUN] Would generate fallback for: {category}")
        return

    client = get_genai_client()
    supabase = get_supabase()
    success = 0
    failed = 0

    for category, prompt_text in CATEGORY_PROMPTS.items():
        try:
            full_prompt = (
                f"Generate a representative photograph for the '{category}' category "
                f"of places in Jeju Aewol area. {prompt_text} "
                f"No text overlays, no watermarks, no people's faces."
            )

            logger.info(f"  Generating fallback for: {category}")
            image_bytes, mime_type = generate_image(client, full_prompt)

            if not image_bytes:
                logger.warning(f"  No image data for {category}")
                failed += 1
                continue

            ext = "webp"
            if "png" in mime_type:
                ext = "png"
            elif "jpeg" in mime_type or "jpg" in mime_type:
                ext = "jpg"

            path = f"fallback/{category}.{ext}"
            url = upload_to_storage(supabase, path, image_bytes, mime_type)

            logger.info(f"  ✅ {category} → {url}")
            success += 1

            # Rate limiting
            time.sleep(2)

        except Exception as e:
            logger.error(f"  ❌ {category} failed: {e}")
            failed += 1

    logger.info(f"\nFallback generation complete: {success} success, {failed} failed")


def generate_spot_images(
    spot_ids: list[str] | None = None,
    category: str | None = None,
    limit: int = 10,
    dry_run: bool = False,
):
    """Generate images for spots"""
    supabase = get_supabase()

    # Build query
    if spot_ids:
        result = supabase.table("story_spots").select(
            "id, name, category"
        ).in_("id", spot_ids).execute()
    elif category:
        result = supabase.table("story_spots").select(
            "id, name, category"
        ).eq("category", category).is_("generated_image_url", "null").limit(limit).execute()
    else:
        result = supabase.table("story_spots").select(
            "id, name, category"
        ).is_("generated_image_url", "null").limit(limit).execute()

    spots = result.data or []
    logger.info(f"Found {len(spots)} spots to generate images for")

    if not spots:
        logger.info("No spots found matching criteria")
        return

    if dry_run:
        for spot in spots:
            logger.info(f"  [DRY-RUN] Would generate: {spot['id']} ({spot['name']}) [{spot['category']}]")
        return

    client = get_genai_client()
    success = 0
    failed = 0

    for i, spot in enumerate(spots, 1):
        try:
            cat = spot.get("category", "coastline")
            base_prompt = CATEGORY_PROMPTS.get(cat, DEFAULT_PROMPT)
            full_prompt = (
                f"Generate a high-quality photograph of {spot['name']} in Jeju Aewol area. "
                f"{base_prompt} "
                f"No text overlays, no watermarks, no people's faces."
            )

            logger.info(f"  [{i}/{len(spots)}] Generating: {spot['name']} ({cat})")
            image_bytes, mime_type = generate_image(client, full_prompt)

            if not image_bytes:
                logger.warning(f"  No image data for {spot['id']}")
                failed += 1
                continue

            ext = "webp"
            if "png" in mime_type:
                ext = "png"
            elif "jpeg" in mime_type or "jpg" in mime_type:
                ext = "jpg"

            path = f"{spot['id']}.{ext}"
            url = upload_to_storage(supabase, path, image_bytes, mime_type)

            # Update DB
            supabase.table("story_spots").update({
                "generated_image_url": url
            }).eq("id", spot["id"]).execute()

            logger.info(f"  ✅ {spot['name']} → {url}")
            success += 1

            # Rate limiting
            time.sleep(2)

        except Exception as e:
            logger.error(f"  ❌ {spot['name']} failed: {e}")
            failed += 1

    logger.info(f"\nSpot image generation complete: {success} success, {failed} failed out of {len(spots)}")


def main():
    parser = argparse.ArgumentParser(description="Batch generate spot images with Gemini Nano Banana")
    parser.add_argument("--fallback", action="store_true", help="Generate category fallback images")
    parser.add_argument("--category", type=str, help="Generate for specific category")
    parser.add_argument("--spot-ids", nargs="+", help="Generate for specific spot IDs")
    parser.add_argument("--limit", type=int, default=10, help="Max spots to generate (default: 10)")
    parser.add_argument("--dry-run", action="store_true", help="Show what would be generated without doing it")

    args = parser.parse_args()

    if args.fallback:
        generate_fallback_images(dry_run=args.dry_run)
    else:
        generate_spot_images(
            spot_ids=args.spot_ids,
            category=args.category,
            limit=args.limit,
            dry_run=args.dry_run,
        )


if __name__ == "__main__":
    main()