"""AI 이미지 생성 스크립트 - GPT Image 1 (OpenAI)과 Gemini Imagen (Google) API를 사용하여
보험 GA 리크루팅 광고 이미지를 생성하고 결과를 JSON 메타데이터로 저장합니다."""

import argparse
import base64
import json
import os
import time
from datetime import datetime
from pathlib import Path
from typing import Any, Optional

import gcloud_auth
import requests

OUTPUT_DIR = Path("/home/jay/workspace/tools/ai-image-gen/output")
ENV_KEYS_FILE = Path("/home/jay/workspace/.env.keys")
RESULTS_JSON = OUTPUT_DIR / "results.json"

COST_TABLE: dict[str, float] = {
    "low": 0.011,
    "medium": 0.042,
    "high": 0.167,
}

SCENARIOS: dict[str, str] = {
    "A": (
        "Professional Korean insurance general agency recruiting advertisement. "
        "Modern office environment with confident business professionals. "
        "Clean, premium design. 1080x1080 square format for Meta ads. "
        "Text overlay: '당신의 커리어, 새로운 시작' in Korean. "
        "Color scheme: navy blue and gold. Corporate premium feel."
    ),
    "B": (
        "Korean insurance consultant personal branding image. "
        "Confident male professional in business suit, modern Seoul office backdrop. "
        "Warm lighting, trustworthy atmosphere. 1080x1080 square. "
        "Subtle text: '보험의 모든 것' in elegant Korean typography."
    ),
    "C": (
        "Motivational career change advertisement for insurance professionals. "
        "Split image: left side dark/stressful office, right side bright/modern workspace. "
        "Aspirational feel. 1080x1080 square format. "
        "Korean text: '지금이 기회입니다' centered."
    ),
}

V2_SCENARIOS: dict[str, str] = {
    "A": (
        "Professional advertisement photograph for Korean insurance general agency recruitment. "
        "Cinematic composition using rule of thirds with strong leading lines. "
        "Environment: luxurious 100th-floor Seoul panoramic office with floor-to-ceiling glass windows revealing city skyline at golden hour, "
        "minimalist interior with warm natural wood accents and a single elegant desk. "
        "Subject: Korean male professional in his early 30s wearing a perfectly tailored navy double-breasted suit with subtle gold cufflinks, "
        "confident subtle smile, chin slightly lifted at 15 degrees, 3/4 angle portrait. "
        "Lighting: natural golden-hour backlight streaming through windows creating beautiful rim lighting on subject's silhouette, "
        "soft diffused fill light illuminating face evenly, subtle volumetric light rays in the air. "
        "Color palette: deep navy (#1B2A4A), warm gold (#C9A84C), soft cream highlights, muted charcoal accents. "
        "Shot on Phase One IQ4 150MP with 80mm lens, f/4 aperture for perfect subject isolation, 8K resolution. "
        "Professional advertising photography quality. Premium aspirational corporate recruitment campaign style. "
        "Bold headline text '새로운 시작' in white sans-serif font, top-center. "
        "No stock photo aesthetic, no clipart, no watermark, no generic corporate imagery. Photorealistic."
    ),
    "B": (
        "Premium personal branding photograph for elite Korean insurance consultant. "
        "Golden ratio composition with shallow depth of field. "
        "Environment: high-end private consultation room with rich dark walnut executive desk, "
        "neatly arranged premium leather portfolio and Mont Blanc pen, warm ambient bookshelf with curated books in soft background bokeh. "
        "Subject: Korean male professional in his 40s wearing charcoal business suit, thin titanium-framed glasses, "
        "warm genuine smile conveying trust and expertise, 3/4 frontal angle, one hand gently resting on desk surface. "
        "Lighting: warm key light from 45 degrees with large softbox creating gentle shadows, "
        "subtle fill light at 1:3 ratio, warm amber ambient glow from background. "
        "Color theory: analogous warm palette — rich amber (#D4A574), cream (#FFF8E7), dark walnut brown (#3C2415), touch of burgundy (#722F37). "
        "Shot on Hasselblad H6D-100c medium format, 100mm f/2.8 lens for creamy bokeh and subject isolation. "
        "Style: premium financial advisor personal branding campaign, like a Rolex or private banking advertisement. "
        "No text overlay — focus purely on image quality and emotional impact. "
        "No stock photo feel, no clipart, no watermark. Photorealistic, warm and trustworthy atmosphere."
    ),
    "C": (
        "Cinematic metaphorical scene for insurance career transition advertisement. "
        "Dramatic wide-angle composition with strong leading lines converging toward bright focal point. "
        "Single powerful scene: silhouette of a Korean professional in business attire walking purposefully "
        "from a dark, cold-toned corporate corridor toward a luminous warm doorway. "
        "Through the doorway: vision of a modern, sunlit co-working space with floor-to-ceiling windows and lush green plants. "
        "Lighting: dramatic volumetric god-rays streaming through the doorway, "
        "high contrast between cool blue-grey shadows (#2C3E50) in the corridor and warm golden light (#F4D03F) from the future. "
        "Subtle lens flare from the bright opening. Floor reflections adding depth. "
        "Color grading: teal-and-orange cinematic palette, cool shadows transitioning to warm highlights. "
        "Shot on ARRI Alexa 65 with Panavision anamorphic 40mm lens, 6K cinema quality, 2.39:1 aspect ratio feel within square frame. "
        "Style: like a Christopher Nolan or Denis Villeneuve film still — epic, contemplative, transformative. "
        "Single line of Korean text '지금, 당신의 차례' in clean white sans-serif, bottom-center, small size. "
        "No stock photo aesthetic, no watermark. Photorealistic, deeply aspirational mood."
    ),
}


def load_env_keys(filepath: Path) -> dict[str, str]:
    """환경변수 파일에서 KEY=VALUE 형식의 키를 로드합니다."""
    keys: dict[str, str] = {}
    if not filepath.exists():
        print(f"[WARNING] 환경변수 파일을 찾을 수 없습니다: {filepath}")
        return keys
    with open(filepath, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            # export KEY=VALUE 또는 KEY=VALUE 형식 모두 처리
            if line.startswith("export "):
                line = line[len("export ") :]
            if "=" in line:
                k, _, v = line.partition("=")
                v = v.strip().strip("\"'")
                keys[k.strip()] = v
    return keys


def load_results(results_file: Path) -> list[dict[str, Any]]:
    """기존 results.json을 로드합니다."""
    if not results_file.exists():
        return []
    with open(results_file, encoding="utf-8") as f:
        return json.load(f)  # type: ignore[no-any-return]


def save_results(results_file: Path, results: list[dict[str, Any]]) -> None:
    """results.json에 결과를 저장합니다."""
    with open(results_file, "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)


def generate_gpt_image(
    api_key: str,
    prompt: str,
    quality: str,
    output_path: Path,
) -> dict[str, Any]:
    """GPT Image 1 API를 호출하여 이미지를 생성하고 저장합니다."""
    url = "https://api.openai.com/v1/images/generations"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    payload = {
        "model": "gpt-image-1",
        "prompt": prompt,
        "n": 1,
        "size": "1024x1024",
        "quality": quality,
    }

    print(f"[GPT] 이미지 생성 요청 중... (quality={quality})")
    start_time = time.time()

    response = requests.post(url, headers=headers, json=payload, timeout=120)
    response.raise_for_status()

    elapsed = time.time() - start_time
    data = response.json()

    # base64 디코딩 후 파일 저장
    image_b64: str = data["data"][0]["b64_json"]
    image_bytes = base64.b64decode(image_b64)
    output_path.write_bytes(image_bytes)

    file_size = output_path.stat().st_size
    print(f"[GPT] 생성 완료: {output_path.name} ({file_size:,} bytes, {elapsed:.1f}초)")

    return {
        "api": "gpt",
        "time_seconds": round(elapsed, 2),
        "file_size_bytes": file_size,
        "cost_usd": COST_TABLE.get(quality, 0.0),
        "error": None,
    }


_GEMINI_SCOPE = "https://www.googleapis.com/auth/generative-language"
_GEMINI_IMAGEN_MODEL = "imagen-3.0-generate-002"


def _generate_gemini_image_with_api_key(
    api_key: str,
    prompt: str,
    output_path: Path,
) -> dict[str, Any]:
    """API 키를 사용하여 Gemini Imagen API로 이미지를 생성합니다."""
    from google import genai

    client = genai.Client(api_key=api_key)

    print(f"[Gemini] API 키 인증으로 이미지 생성 요청 중... (모델: {_GEMINI_IMAGEN_MODEL})")
    start_time = time.time()

    response = client.models.generate_images(
        model=_GEMINI_IMAGEN_MODEL,
        prompt=prompt,
        config=genai.types.GenerateImagesConfig(
            number_of_images=1,
            output_mime_type="image/png",
        ),
    )

    elapsed = time.time() - start_time

    if not response.generated_images:
        raise RuntimeError("Gemini API가 이미지를 반환하지 않았습니다.")
    generated = response.generated_images[0]
    if generated.image is None:
        raise RuntimeError("생성된 이미지 데이터가 없습니다.")
    generated.image.save(str(output_path))

    file_size = output_path.stat().st_size
    print(f"[Gemini] 생성 완료: {output_path.name} ({file_size:,} bytes, {elapsed:.1f}초)")

    return {
        "api": "gemini",
        "auth_method": "api_key",
        "time_seconds": round(elapsed, 2),
        "file_size_bytes": file_size,
        "cost_usd": None,
        "error": None,
    }


def _generate_gemini_image_with_sa_token(
    prompt: str,
    output_path: Path,
) -> dict[str, Any]:
    """SA Bearer 토큰을 사용하여 Gemini generateContent REST API로 이미지를 생성합니다."""
    token = gcloud_auth.get_service_account_token(_GEMINI_SCOPE)

    gemini_api_base = "https://generativelanguage.googleapis.com/v1beta"
    # Imagen은 generateImages 엔드포인트를 사용하지만, Bearer 토큰은 generateContent에서 지원됨
    # flash 모델로 fallback하여 generateContent 방식 사용
    model_id = "gemini-3.1-flash-image-preview"
    url = f"{gemini_api_base}/models/{model_id}:generateContent"
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }
    payload: dict[str, Any] = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {"responseModalities": ["IMAGE", "TEXT"]},
    }

    print(f"[Gemini] SA 토큰 fallback으로 이미지 생성 요청 중... (모델: {model_id})")
    start_time = time.time()
    response = requests.post(url, headers=headers, json=payload, timeout=180)
    response.raise_for_status()
    elapsed = time.time() - start_time

    data: dict[str, Any] = response.json()
    candidates = data.get("candidates", [])
    if not candidates:
        raise RuntimeError(f"SA fallback 응답에 candidates가 없습니다.")

    parts = candidates[0].get("content", {}).get("parts", [])
    image_part: dict[str, Any] | None = None
    for part in parts:
        if "inlineData" in part:
            image_part = part
            break

    if image_part is None:
        raise RuntimeError("SA fallback 응답에 이미지 데이터가 없습니다.")

    mime_type: str = image_part["inlineData"].get("mimeType", "image/jpeg")
    image_bytes = base64.b64decode(image_part["inlineData"]["data"])
    ext = ".jpg" if "jpeg" in mime_type else ".png"
    if output_path.suffix.lower() != ext:
        output_path = output_path.with_suffix(ext)
    output_path.write_bytes(image_bytes)

    file_size = output_path.stat().st_size
    print(f"[Gemini] SA fallback 생성 완료: {output_path.name} ({file_size:,} bytes, {elapsed:.1f}초)")

    return {
        "api": "gemini",
        "auth_method": "sa_bearer_token_fallback",
        "time_seconds": round(elapsed, 2),
        "file_size_bytes": file_size,
        "cost_usd": None,
        "error": None,
    }


def generate_gemini_image(
    api_key: str,
    prompt: str,
    output_path: Path,
) -> dict[str, Any]:
    """Gemini 이미지를 생성합니다.

    인증 우선순위:
    1. GEMINI_API_KEY가 있으면 API 키 인증 (기존 방식)
    2. API 키 실패(429 할당량 초과 등) 시 SA Bearer 토큰으로 fallback

    Args:
        api_key: Gemini API 키 (빈 문자열이면 SA 토큰으로 직행).
        prompt: 이미지 생성 프롬프트.
        output_path: 저장할 파일 경로.

    Returns:
        생성 결과 메타데이터 딕셔너리.
    """
    if api_key:
        try:
            return _generate_gemini_image_with_api_key(api_key, prompt, output_path)
        except requests.HTTPError as e:
            status = e.response.status_code if e.response is not None else 0
            if status == 429:
                print(f"[Gemini] API 키 할당량 초과 (429). SA Bearer 토큰 fallback 시도 중...")
            else:
                print(f"[Gemini] API 키 인증 HTTP 오류 ({status}). SA Bearer 토큰 fallback 시도 중...")
        except Exception as e:
            print(f"[Gemini] API 키 인증 실패: {type(e).__name__}: {e}. SA Bearer 토큰 fallback 시도 중...")

    # SA Bearer 토큰 fallback
    return _generate_gemini_image_with_sa_token(prompt, output_path)


def run_generation(
    api: str,
    scenario: str,
    quality: str,
    api_keys: dict[str, str],
    output_dir: Optional[Path] = None,
    prompt_version: str = "v1",
) -> dict[str, Any]:
    """단일 이미지 생성 작업을 실행합니다."""
    if output_dir is None:
        output_dir = OUTPUT_DIR

    scenarios_map = V2_SCENARIOS if prompt_version == "v2" else SCENARIOS
    prompt = scenarios_map[scenario]
    filename = f"{api}_{scenario}_{quality}.png"
    output_path = output_dir / filename
    timestamp = datetime.now().isoformat()

    result: dict[str, Any] = {
        "api": api,
        "scenario": scenario,
        "quality": quality,
        "prompt_version": prompt_version,
        "filename": filename,
        "timestamp": timestamp,
        "time_seconds": None,
        "file_size_bytes": None,
        "cost_usd": None,
        "error": None,
    }

    try:
        if api == "gpt":
            api_key = api_keys.get("OPENAI_API_KEY", "")
            if not api_key:
                print("[WARNING] OPENAI_API_KEY가 없어 GPT API를 스킵합니다.")
                result["error"] = "OPENAI_API_KEY not found"
                return result
            gen_result = generate_gpt_image(api_key, prompt, quality, output_path)

        elif api == "gemini":
            api_key = api_keys.get("GEMINI_API_KEY", "")
            if not api_key:
                print("[WARNING] GEMINI_API_KEY 없음. SA Bearer 토큰 fallback으로 시도합니다.")
            gen_result = generate_gemini_image(api_key, prompt, output_path)

        else:
            result["error"] = f"Unknown API: {api}"
            return result

        result.update(gen_result)

    except requests.HTTPError as e:
        error_msg = f"HTTP 오류: {e.response.status_code} - {e.response.text[:200]}"
        print(f"[ERROR] {error_msg}")
        result["error"] = error_msg
    except Exception as e:
        error_msg = f"{type(e).__name__}: {e}"
        print(f"[ERROR] {error_msg}")
        result["error"] = error_msg

    return result


def parse_args() -> argparse.Namespace:
    """CLI 인자를 파싱합니다."""
    parser = argparse.ArgumentParser(
        description="GPT Image 1 / Gemini Imagen을 사용한 보험 GA 광고 이미지 생성",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
예시:
  python3 generate_ai_image.py --api gpt --scenario A --quality high
  python3 generate_ai_image.py --api gemini --scenario B --quality medium
  python3 generate_ai_image.py --api gpt --scenario A --quality high --prompt-version v2
  python3 generate_ai_image.py --all
  python3 generate_ai_image.py --all --prompt-version v2
        """,
    )
    parser.add_argument(
        "--api",
        choices=["gpt", "gemini"],
        help="사용할 API (gpt 또는 gemini)",
    )
    parser.add_argument(
        "--scenario",
        choices=["A", "B", "C"],
        help="시나리오 선택 (A: GA리크루팅, B: 전문성브랜딩, C: 감성전직유도)",
    )
    parser.add_argument(
        "--quality",
        choices=["low", "medium", "high"],
        default="medium",
        help="이미지 품질 (기본값: medium, GPT 전용 파라미터)",
    )
    parser.add_argument(
        "--all",
        action="store_true",
        dest="run_all",
        help="모든 시나리오를 사용 가능한 모든 API로 실행",
    )
    parser.add_argument(
        "--prompt-version",
        choices=["v1", "v2"],
        default="v1",
        dest="prompt_version",
        help="프롬프트 버전 선택 (기본값: v1)",
    )
    return parser.parse_args()


def get_available_apis(api_keys: dict[str, str]) -> list[str]:
    """사용 가능한 API 목록을 반환합니다."""
    available: list[str] = []
    if api_keys.get("OPENAI_API_KEY"):
        available.append("gpt")
    else:
        print("[WARNING] OPENAI_API_KEY가 없어 GPT API를 사용할 수 없습니다.")
    if api_keys.get("GEMINI_API_KEY"):
        available.append("gemini")
    else:
        print("[WARNING] GEMINI_API_KEY 없음. SA Bearer 토큰 fallback을 시도합니다.")
        available.append("gemini")
    return available


def main() -> None:
    """메인 실행 함수."""
    args = parse_args()

    # --all 없이 --api, --scenario 미지정 시 오류
    if not args.run_all and (args.api is None or args.scenario is None):
        print("[ERROR] --all 옵션이 없으면 --api와 --scenario를 모두 지정해야 합니다.")
        raise SystemExit(1)

    # 프롬프트 버전에 따른 출력 디렉토리 결정
    if args.prompt_version == "v2":
        active_output_dir = OUTPUT_DIR / "v2-gpt-advanced"
    else:
        active_output_dir = OUTPUT_DIR / "v1-gpt-poc"

    active_output_dir.mkdir(parents=True, exist_ok=True)
    active_results_json = active_output_dir / "results.json"

    # API 키 로드
    api_keys = load_env_keys(ENV_KEYS_FILE)
    # 환경변수에서도 읽기 (환경변수가 파일보다 우선)
    for key in ("OPENAI_API_KEY", "GEMINI_API_KEY"):
        env_val = os.environ.get(key)
        if env_val:
            api_keys[key] = env_val

    # 실행할 작업 목록 구성
    tasks: list[tuple[str, str, str]] = []

    if args.run_all:
        available_apis = get_available_apis(api_keys)
        if not available_apis:
            print("[ERROR] 사용 가능한 API가 없습니다.")
            raise SystemExit(1)
        for api in available_apis:
            for scenario in ["A", "B", "C"]:
                quality = args.quality if api == "gpt" else "medium"
                tasks.append((api, scenario, quality))
    else:
        tasks.append((args.api, args.scenario, args.quality))

    print(f"\n총 {len(tasks)}개 이미지 생성 작업을 시작합니다. (프롬프트 버전: {args.prompt_version})\n")

    results = load_results(active_results_json)

    for i, (api, scenario, quality) in enumerate(tasks, 1):
        print(f"[{i}/{len(tasks)}] API={api}, 시나리오={scenario}, 품질={quality}")
        result = run_generation(
            api,
            scenario,
            quality,
            api_keys,
            output_dir=active_output_dir,
            prompt_version=args.prompt_version,
        )
        results.append(result)
        save_results(active_results_json, results)
        print(f"  -> 결과 저장: {result.get('filename', 'N/A')} " f"(오류: {result.get('error', '없음')})\n")

    print(f"모든 작업 완료. 결과 저장: {active_results_json}")


if __name__ == "__main__":
    main()