#!/usr/bin/env python3
"""컨셉 #18 배경 이미지 생성 - Gemini API"""

import base64
import sys
from pathlib import Path

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

OUTPUT_DIR = Path("/home/jay/workspace/output/meta-ads/concept-catalog/18-portrait-text")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

PROMPT = (
    "Professional Korean business person in their 30s wearing a crisp white dress shirt, "
    "confident warm smile, clean modern office background with soft natural lighting, "
    "upper body portrait shot, high quality professional headshot photography style, "
    "the person is positioned on the right side of the frame leaving the left half mostly empty "
    "or with a solid dark navy (#1A2E4A) background, 1080x1080 pixels"
)

GEMINI_API_BASE = "https://generativelanguage.googleapis.com/v1beta"
SCOPE = "https://www.googleapis.com/auth/generative-language"

IMAGE_MODELS = [
    "gemini-2.5-flash-image",
    "gemini-3-pro-image-preview",
    "gemini-3.1-flash-image-preview",
]


def generate_with_gemini(token: str, model_id: str) -> bytes | None:
    """Gemini generateContent API로 이미지 생성"""
    url = f"{GEMINI_API_BASE}/models/{model_id}:generateContent"
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }
    payload = {
        "contents": [{"parts": [{"text": PROMPT}]}],
        "generationConfig": {"responseModalities": ["IMAGE", "TEXT"]},
    }
    print(f"[배경생성] 모델: {model_id}")
    resp = requests.post(url, headers=headers, json=payload, timeout=180)
    print(f"  HTTP 상태: {resp.status_code}")

    if resp.status_code != 200:
        print(f"  오류: {resp.text[:300]}")
        return None

    data = resp.json()
    candidates = data.get("candidates", [])
    if not candidates:
        print("  candidates 없음")
        return None

    parts = candidates[0].get("content", {}).get("parts", [])
    for part in parts:
        if "inlineData" in part:
            return base64.b64decode(part["inlineData"]["data"])

    print("  이미지 데이터 없음")
    return None


def generate_with_imagen4(token: str) -> bytes | None:
    """Imagen 4 사용 (최후 fallback)"""
    url = f"{GEMINI_API_BASE}/models/imagen-4.0-generate-001:predict"
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }
    payload = {
        "instances": [{"prompt": PROMPT}],
        "parameters": {
            "sampleCount": 1,
            "aspectRatio": "1:1",
        }
    }
    print("[배경생성] fallback 모델: imagen-4.0-generate-001")
    resp = requests.post(url, headers=headers, json=payload, timeout=180)
    print(f"  HTTP 상태: {resp.status_code}")

    if resp.status_code != 200:
        print(f"  오류: {resp.text[:300]}")
        return None

    data = resp.json()
    predictions = data.get("predictions", [])
    if not predictions:
        print("  predictions 없음")
        return None

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

    return base64.b64decode(b64)


def main():
    print("[인증] SA 토큰 획득 중...")
    token = gcloud_auth.get_service_account_token(SCOPE)
    print(f"[인증] 토큰 획득 성공 (길이: {len(token)})")

    # Gemini 이미지 모델 순서대로 시도
    image_bytes = None
    for model in IMAGE_MODELS:
        image_bytes = generate_with_gemini(token, model)
        if image_bytes is not None:
            break

    # 실패 시 Imagen 4 시도
    if image_bytes is None:
        image_bytes = generate_with_imagen4(token)

    if image_bytes is None:
        print("[오류] 모든 모델 실패")
        sys.exit(1)

    # 저장
    out_path = OUTPUT_DIR / "bg.jpg"
    out_path.write_bytes(image_bytes)
    print(f"\n[완료] 배경 이미지 저장: {out_path} ({len(image_bytes):,} bytes)")


if __name__ == "__main__":
    main()
