#!/usr/bin/env python3
"""Cell-4 GA × 정당한 보상 배너 배경 이미지 생성 스크립트"""

import base64
import os
import sys
from pathlib import Path

sys.path.insert(0, "/home/jay/workspace/tools/ai-image-gen")

OUTPUT_DIR = Path("/home/jay/workspace/output/banners/cell-4-ga-fair")
FALLBACK_BG = Path("/home/jay/workspace/output/meta-ads/concept-catalog/48-madstars-digital-interactive/bg.jpg")

PROMPTS = {
    "bg-cell4-1080": (
        "Dramatic aerial view of a crossroads intersection at night in a Korean city, "
        "red traffic lights reflecting on wet asphalt, moody dark atmosphere with red and black tones, "
        "cinematic photography style, photorealistic urban scene, no people, no text, no logos. "
        "Square 1:1 format."
    ),
    "bg-cell4-1200": (
        "Wide cinematic view of a dark Korean financial district at night, "
        "red neon signs reflecting on glass buildings, urgent and dramatic atmosphere, "
        "rain-wet streets, photorealistic, no people, no text, no logos. "
        "Wide landscape 1.91:1 format."
    ),
}


def generate_with_api_key(prompt: str, output_path: Path) -> bool:
    """Gemini REST API (API 키)로 이미지 생성"""
    import requests

    api_key = os.environ.get("GEMINI_API_KEY", "AIzaSyDEy7IcOoTbQsI4nxfPOw3ZFbYqEL_PgFU")

    url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-3-pro-image-preview:generateContent?key={api_key}"

    payload = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {"responseModalities": ["TEXT", "IMAGE"]},
    }

    print(f"  Gemini API 호출 중: {output_path.name} ...")
    try:
        resp = requests.post(url, json=payload, timeout=90)
        resp.raise_for_status()
        data = resp.json()

        candidates = data.get("candidates", [])
        if not candidates:
            print(f"  [WARN] candidates 없음: {data}")
            return False

        for part in candidates[0].get("content", {}).get("parts", []):
            if "inlineData" in part:
                img_bytes = base64.b64decode(part["inlineData"]["data"])
                output_path.write_bytes(img_bytes)
                print(f"  [OK] 저장: {output_path} ({len(img_bytes)//1024}KB)")
                return True

        print(f"  [WARN] 이미지 파트 없음")
        return False

    except Exception as e:
        print(f"  [ERROR] API 호출 실패: {e}")
        return False


def try_imagen(prompt: str, output_path: Path) -> bool:
    """Imagen 3 API로 이미지 생성 시도"""
    import requests

    api_key = os.environ.get("GEMINI_API_KEY", "AIzaSyDEy7IcOoTbQsI4nxfPOw3ZFbYqEL_PgFU")

    url = f"https://generativelanguage.googleapis.com/v1beta/models/imagen-4.0-generate-001:predict?key={api_key}"

    aspect = "1:1" if "1080" in output_path.name else "16:9"

    payload = {
        "instances": [{"prompt": prompt}],
        "parameters": {
            "sampleCount": 1,
            "aspectRatio": aspect,
            "safetyFilterLevel": "block_few",
            "personGeneration": "dont_allow",
        },
    }

    print(f"  Imagen 3 API 호출 중: {output_path.name} ({aspect}) ...")
    try:
        resp = requests.post(url, json=payload, timeout=90)
        resp.raise_for_status()
        data = resp.json()

        predictions = data.get("predictions", [])
        if not predictions:
            print(f"  [WARN] predictions 없음: {data}")
            return False

        b64 = predictions[0].get("bytesBase64Encoded", "")
        if not b64:
            print(f"  [WARN] 이미지 데이터 없음")
            return False

        img_bytes = base64.b64decode(b64)
        output_path.write_bytes(img_bytes)
        print(f"  [OK] 저장: {output_path} ({len(img_bytes)//1024}KB)")
        return True

    except Exception as e:
        print(f"  [ERROR] Imagen API 호출 실패: {e}")
        return False


def use_fallback(output_path: Path) -> bool:
    """폴백 이미지 복사"""
    import shutil
    if FALLBACK_BG.exists():
        shutil.copy2(FALLBACK_BG, output_path)
        print(f"  [FALLBACK] {FALLBACK_BG} → {output_path}")
        return True
    print(f"  [ERROR] 폴백 파일도 없음: {FALLBACK_BG}")
    return False


def main():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    results = {}

    for name, prompt in PROMPTS.items():
        out = OUTPUT_DIR / f"{name}.png"
        print(f"\n[{name}] 배경 생성 시작")

        # 1차: Imagen 3 시도
        ok = try_imagen(prompt, out)

        # 2차: Gemini Flash Image Generation 시도
        if not ok:
            ok = generate_with_api_key(prompt, out)

        # 3차: 폴백
        if not ok:
            ok = use_fallback(out)

        results[name] = "OK" if ok else "FAILED"

    print("\n=== 배경 생성 결과 ===")
    for name, status in results.items():
        print(f"  {name}: {status}")

    return 0 if all(v == "OK" for v in results.values()) else 1


if __name__ == "__main__":
    sys.exit(main())