"""Gemini API를 사용한 v3 캠페인 이미지 생성 스크립트.

gemini-2.0-flash-preview-image-generation 모델을 통해
v3_prompts.json에 정의된 프롬프트 기반으로 이미지를 생성합니다.
한글 깨짐 방지를 위한 다중 시도(attempt) 및 결과 기록 기능을 지원합니다.
"""

from __future__ import annotations

import argparse
import base64
import json
import logging
import shutil
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any

import requests

# ──────────────────────────────────────────────
# 상수 정의
# ──────────────────────────────────────────────
MODEL_ID = "gemini-3-pro-image-preview"
GEMINI_API_KEY = "AIzaSyB4X1239KxqIYIuhr9rBrawwxs2RddbYcc"
GEMINI_API_ENDPOINT = (
    "https://generativelanguage.googleapis.com/v1beta/models/"
    f"{MODEL_ID}:generateContent?key={GEMINI_API_KEY}"
)

PROMPTS_FILE = Path("/home/jay/workspace/teams/dev1/v3_prompts.json")

OUTPUT_BASE = Path("/home/jay/workspace/output/campaign-top/v3-gemini")
OUTPUT_DIRS: dict[str, Path] = {
    "meta_carousel_a": OUTPUT_BASE / "meta-carousel-a",
    "meta_carousel_b": OUTPUT_BASE / "meta-carousel-b",
    "naver_gfa": OUTPUT_BASE / "naver-gfa",
}

RESULTS_JSON = OUTPUT_BASE / "generation_results.json"

# 그룹 선택 옵션 → 프롬프트 JSON 키 매핑
GROUP_MAP: dict[str, list[str]] = {
    "all": ["meta_carousel_a", "meta_carousel_b", "naver_gfa"],
    "a": ["meta_carousel_a"],
    "b": ["meta_carousel_b"],
    "gfa": ["naver_gfa"],
}

REQUEST_TIMEOUT = 300  # 초
INTER_IMAGE_DELAY = 5  # 이미지 간 대기 시간(초, Pro 모델 rate limit 대응)

# ──────────────────────────────────────────────
# 로거 설정
# ──────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("v3_generate")


# ──────────────────────────────────────────────
# 인증
# ──────────────────────────────────────────────
def get_auth_token() -> str:
    """API 키를 반환합니다. (하위호환용 — API 키가 URL에 포함되므로 빈 문자열 반환)"""
    logger.info("API 키 인증 사용 (URL 파라미터)")
    return ""


# ──────────────────────────────────────────────
# 프롬프트 파일 읽기
# ──────────────────────────────────────────────
def load_prompts(prompts_file: Path) -> dict[str, list[dict[str, Any]]]:
    """v3_prompts.json 파일에서 프롬프트 설정을 읽어옵니다.

    Args:
        prompts_file: 프롬프트 JSON 파일 경로.

    Returns:
        그룹명 → 슬라이드 목록 딕셔너리.

    Raises:
        SystemExit: 파일이 없거나 JSON 파싱 실패 시.
    """
    if not prompts_file.exists():
        logger.error("프롬프트 파일을 찾을 수 없습니다: %s", prompts_file)
        raise SystemExit(1)

    try:
        with open(prompts_file, encoding="utf-8") as f:
            data: dict[str, list[dict[str, Any]]] = json.load(f)
        logger.info("프롬프트 파일 로드 완료: %s", prompts_file)
        return data
    except json.JSONDecodeError as e:
        logger.error("프롬프트 파일 JSON 파싱 실패: %s", e)
        raise SystemExit(1) from e


# ──────────────────────────────────────────────
# Gemini API 호출
# ──────────────────────────────────────────────
def call_gemini_api(
    token: str,
    prompt: str,
) -> dict[str, Any]:
    """Gemini API에 이미지 생성 요청을 보냅니다.

    Args:
        token: Bearer 인증 토큰.
        prompt: 이미지 생성 프롬프트.

    Returns:
        API 응답 JSON 딕셔너리.

    Raises:
        requests.HTTPError: HTTP 4xx/5xx 오류 발생 시.
        requests.RequestException: 네트워크 오류 발생 시.
    """
    headers = {
        "Content-Type": "application/json",
    }
    payload: dict[str, Any] = {
        "contents": [{"parts": [{"text": prompt}]}],
        "generationConfig": {
            "responseModalities": ["IMAGE", "TEXT"],
            "temperature": 1,
            "topP": 0.95,
            "topK": 40,
        },
    }

    logger.debug("API 호출: %s", GEMINI_API_ENDPOINT)
    response = requests.post(
        GEMINI_API_ENDPOINT,
        headers=headers,
        json=payload,
        timeout=REQUEST_TIMEOUT,
    )

    if not response.ok:
        logger.error(
            "HTTP 오류 %d: %s",
            response.status_code,
            response.text[:500],
        )
        response.raise_for_status()

    return response.json()  # type: ignore[no-any-return]


# ──────────────────────────────────────────────
# 응답 파싱
# ──────────────────────────────────────────────
def parse_image_from_response(
    response_data: dict[str, Any],
) -> tuple[bytes, str]:
    """Gemini API 응답에서 이미지 데이터와 MIME 타입을 추출합니다.

    Args:
        response_data: Gemini API 응답 JSON.

    Returns:
        (이미지 바이트, mime_type) 튜플.

    Raises:
        ValueError: 이미지 데이터가 응답에 없는 경우.
    """
    candidates = response_data.get("candidates", [])
    if not candidates:
        raise ValueError(
            f"응답에 candidates가 없습니다. 응답: {json.dumps(response_data)[:300]}"
        )

    parts = candidates[0].get("content", {}).get("parts", [])

    # inlineData가 있는 파트 탐색
    for part in parts:
        if "inlineData" in part:
            mime_type: str = part["inlineData"].get("mimeType", "image/jpeg")
            image_b64: str = part["inlineData"]["data"]
            image_bytes = base64.b64decode(image_b64)
            return image_bytes, mime_type

    # 이미지가 없는 경우 텍스트 파트 로깅
    text_parts = [p.get("text", "") for p in parts if "text" in p]
    if text_parts:
        logger.warning("이미지 데이터 없음. 텍스트 응답: %s", text_parts[:2])
    raise ValueError(
        f"이미지 데이터가 응답에 없습니다. 텍스트 파트: {text_parts[:2]}"
    )


# ──────────────────────────────────────────────
# 이미지 저장
# ──────────────────────────────────────────────
def save_image(
    image_bytes: bytes,
    mime_type: str,
    output_path: Path,
) -> Path:
    """이미지 바이트를 MIME 타입에 맞는 확장자로 저장합니다.

    Args:
        image_bytes: 저장할 이미지 바이트 데이터.
        mime_type: 이미지 MIME 타입 (예: image/jpeg, image/png).
        output_path: 저장 경로 (확장자는 MIME 타입에 따라 보정됨).

    Returns:
        실제 저장된 파일 경로.
    """
    ext = ".jpg" if "jpeg" in mime_type else ".png"
    if output_path.suffix.lower() != ext:
        output_path = output_path.with_suffix(ext)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_bytes(image_bytes)
    logger.info(
        "이미지 저장 완료: %s (%s bytes, mime=%s)",
        output_path.name,
        f"{output_path.stat().st_size:,}",
        mime_type,
    )
    return output_path


# ──────────────────────────────────────────────
# 단일 이미지 생성 (attempt 포함)
# ──────────────────────────────────────────────
def generate_single_image(
    token: str,
    prompt: str,
    base_output_path: Path,
    attempt: int,
) -> dict[str, Any]:
    """단일 attempt에 대해 이미지를 생성하고 저장합니다.

    파일명 규칙: `{stem}_attempt{attempt}{ext}`

    Args:
        token: Bearer 인증 토큰.
        prompt: 이미지 생성 프롬프트.
        base_output_path: 기본 출력 경로 (확장자 포함).
        attempt: 현재 시도 번호 (1부터 시작).

    Returns:
        생성 결과 딕셔너리 (success, file_path, file_size, elapsed, error 등).
    """
    stem = base_output_path.stem
    suffix = base_output_path.suffix or ".png"
    attempt_path = base_output_path.with_name(f"{stem}_attempt{attempt}{suffix}")

    result: dict[str, Any] = {
        "attempt": attempt,
        "file_path": str(attempt_path),
        "file_size": None,
        "elapsed_seconds": None,
        "success": False,
        "error": None,
    }

    start_time = time.time()
    try:
        response_data = call_gemini_api(token, prompt)
        image_bytes, mime_type = parse_image_from_response(response_data)
        saved_path = save_image(image_bytes, mime_type, attempt_path)

        elapsed = time.time() - start_time
        result.update(
            {
                "attempt": attempt,
                "file_path": str(saved_path),
                "file_size": saved_path.stat().st_size,
                "elapsed_seconds": round(elapsed, 2),
                "mime_type": mime_type,
                "success": True,
                "error": None,
            }
        )
    except requests.HTTPError as e:
        elapsed = time.time() - start_time
        error_msg = (
            f"HTTP {e.response.status_code}: {e.response.text[:500]}"
        )
        logger.error("Attempt %d HTTP 오류: %s", attempt, error_msg)
        result["elapsed_seconds"] = round(elapsed, 2)
        result["error"] = error_msg
    except ValueError as e:
        elapsed = time.time() - start_time
        logger.error("Attempt %d 이미지 파싱 오류: %s", attempt, e)
        result["elapsed_seconds"] = round(elapsed, 2)
        result["error"] = str(e)
    except Exception as e:
        elapsed = time.time() - start_time
        error_msg = f"{type(e).__name__}: {e}"
        logger.error("Attempt %d 예기치 않은 오류: %s", attempt, error_msg)
        result["elapsed_seconds"] = round(elapsed, 2)
        result["error"] = error_msg

    return result


# ──────────────────────────────────────────────
# 슬라이드 처리 (재시도 포함)
# ──────────────────────────────────────────────
def process_slide(
    token: str,
    group_key: str,
    slide_config: dict[str, Any],
    max_retries: int,
    dry_run: bool,
) -> dict[str, Any]:
    """슬라이드 하나에 대해 최대 max_retries번 이미지를 생성합니다.

    성공한 마지막 attempt 파일을 기본 파일명(output_file)으로도 복사합니다.

    Args:
        token: Bearer 인증 토큰.
        group_key: 프롬프트 그룹 키 (예: meta_carousel_a).
        slide_config: 슬라이드 설정 딕셔너리.
        max_retries: 최대 재시도 횟수.
        dry_run: True이면 API 호출 없이 프롬프트만 출력.

    Returns:
        슬라이드 전체 결과 딕셔너리.
    """
    slide_num = slide_config.get("slide", "?")
    name = slide_config.get("name", "unknown")
    prompt: str = slide_config.get("prompt", "")
    output_file: str = slide_config.get("output_file", f"slide_{slide_num}_{name}.png")

    output_dir = OUTPUT_DIRS[group_key]
    base_output_path = output_dir / output_file

    prompt_summary = prompt[:50]
    logger.info(
        "[%s] 슬라이드 %s (%s) 처리 시작 | 프롬프트: %s...",
        group_key,
        slide_num,
        name,
        prompt_summary,
    )

    slide_result: dict[str, Any] = {
        "group": group_key,
        "slide": slide_num,
        "name": name,
        "prompt_summary": prompt_summary,
        "output_file": str(base_output_path),
        "attempts": [],
        "total_attempts": 0,
        "success": False,
        "final_file_path": None,
        "final_file_size": None,
        "total_elapsed_seconds": None,
        "timestamp": datetime.now().isoformat(),
        "error": None,
    }

    if dry_run:
        print(f"\n[DRY-RUN] 그룹: {group_key} | 슬라이드: {slide_num} | 이름: {name}")
        print(f"  출력 경로: {base_output_path}")
        print(f"  프롬프트: {prompt}")
        slide_result["success"] = True
        slide_result["error"] = "dry-run (API 호출 생략)"
        return slide_result

    output_dir.mkdir(parents=True, exist_ok=True)

    total_start = time.time()
    last_success_path: Path | None = None

    for attempt in range(1, max_retries + 1):
        logger.info(
            "[%s] 슬라이드 %s attempt %d/%d 시작",
            group_key,
            slide_num,
            attempt,
            max_retries,
        )

        attempt_result = generate_single_image(token, prompt, base_output_path, attempt)
        slide_result["attempts"].append(attempt_result)

        if attempt_result["success"]:
            last_success_path = Path(attempt_result["file_path"])
            logger.info(
                "[%s] 슬라이드 %s attempt %d 성공: %s",
                group_key,
                slide_num,
                attempt,
                last_success_path.name,
            )
        else:
            logger.warning(
                "[%s] 슬라이드 %s attempt %d 실패: %s",
                group_key,
                slide_num,
                attempt,
                attempt_result["error"],
            )

        # 마지막 attempt가 아니면 rate limit 방지 대기
        if attempt < max_retries:
            logger.debug("다음 attempt 전 %d초 대기...", INTER_IMAGE_DELAY)
            time.sleep(INTER_IMAGE_DELAY)

    total_elapsed = time.time() - total_start
    slide_result["total_attempts"] = max_retries
    slide_result["total_elapsed_seconds"] = round(total_elapsed, 2)

    if last_success_path is not None:
        # 마지막 성공 attempt를 기본 파일명으로 복사
        final_path = output_dir / output_file
        # 확장자 보정: last_success_path 확장자 기준
        if last_success_path.suffix != final_path.suffix:
            final_path = final_path.with_suffix(last_success_path.suffix)
        shutil.copy2(last_success_path, final_path)
        logger.info(
            "[%s] 슬라이드 %s 최종 파일 복사: %s → %s",
            group_key,
            slide_num,
            last_success_path.name,
            final_path.name,
        )

        slide_result["success"] = True
        slide_result["final_file_path"] = str(final_path)
        slide_result["final_file_size"] = final_path.stat().st_size
        slide_result["error"] = None
    else:
        last_error = slide_result["attempts"][-1]["error"] if slide_result["attempts"] else "알 수 없는 오류"
        slide_result["success"] = False
        slide_result["error"] = f"모든 {max_retries}회 시도 실패. 마지막 오류: {last_error}"
        logger.error(
            "[%s] 슬라이드 %s 모든 시도 실패: %s",
            group_key,
            slide_num,
            slide_result["error"],
        )

    return slide_result


# ──────────────────────────────────────────────
# 결과 저장
# ──────────────────────────────────────────────
def save_results(
    results: list[dict[str, Any]],
    total_elapsed: float,
) -> None:
    """생성 결과를 generation_results.json에 저장합니다.

    Args:
        results: 슬라이드별 결과 리스트.
        total_elapsed: 전체 소요 시간(초).
    """
    RESULTS_JSON.parent.mkdir(parents=True, exist_ok=True)

    success_count = sum(1 for r in results if r.get("success"))
    fail_count = len(results) - success_count

    output_data: dict[str, Any] = {
        "run_timestamp": datetime.now().isoformat(),
        "model": MODEL_ID,
        "total_slides": len(results),
        "success_count": success_count,
        "fail_count": fail_count,
        "total_elapsed_seconds": round(total_elapsed, 2),
        "results": results,
    }

    with open(RESULTS_JSON, "w", encoding="utf-8") as f:
        json.dump(output_data, f, ensure_ascii=False, indent=2)

    logger.info("결과 저장 완료: %s", RESULTS_JSON)


# ──────────────────────────────────────────────
# CLI 인터페이스
# ──────────────────────────────────────────────
def parse_args() -> argparse.Namespace:
    """CLI 인자를 파싱합니다.

    Returns:
        파싱된 argparse.Namespace 객체.
    """
    parser = argparse.ArgumentParser(
        description="Gemini API v3 캠페인 이미지 생성 스크립트",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
예시:
  python v3_generate.py                    # 전체 생성
  python v3_generate.py --group a          # meta_carousel_a만 생성
  python v3_generate.py --group gfa        # naver_gfa만 생성
  python v3_generate.py --max-retries 5    # 최대 5회 재시도
  python v3_generate.py --dry-run          # API 호출 없이 프롬프트 확인
""",
    )
    parser.add_argument(
        "--group",
        choices=["all", "a", "b", "gfa"],
        default="all",
        help="생성할 그룹 선택 (기본: all)",
    )
    parser.add_argument(
        "--max-retries",
        type=int,
        default=3,
        metavar="N",
        help="이미지당 최대 재시도 횟수 (기본: 3)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="API 호출 없이 프롬프트 및 설정만 출력",
    )
    return parser.parse_args()


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

    print("=" * 70)
    print("v3 캠페인 이미지 생성 시작")
    print(f"모델: {MODEL_ID}")
    print(f"그룹: {args.group}")
    print(f"최대 재시도: {args.max_retries}회")
    print(f"Dry-run: {args.dry_run}")
    print(f"출력 디렉토리: {OUTPUT_BASE}")
    print("=" * 70)

    # 프롬프트 파일 로드
    prompts = load_prompts(PROMPTS_FILE)

    # 처리할 그룹 목록 결정
    target_groups = GROUP_MAP[args.group]
    logger.info("처리 대상 그룹: %s", target_groups)

    # 인증 토큰 획득 (dry-run이 아닐 경우)
    token: str = ""
    if not args.dry_run:
        token = get_auth_token()

    all_results: list[dict[str, Any]] = []
    total_start = time.time()
    slide_index = 0

    for group_key in target_groups:
        if group_key not in prompts:
            logger.warning("프롬프트 파일에 그룹 '%s'가 없습니다. 건너뜁니다.", group_key)
            continue

        slides = prompts[group_key]
        logger.info(
            "\n[%s] 그룹 처리 시작 (%d개 슬라이드)", group_key, len(slides)
        )

        for slide_config in slides:
            # 첫 번째 슬라이드가 아닌 경우 이미지 간 대기
            if slide_index > 0 and not args.dry_run:
                logger.debug("rate limit 방지: %d초 대기 중...", INTER_IMAGE_DELAY)
                time.sleep(INTER_IMAGE_DELAY)

            slide_result = process_slide(
                token=token,
                group_key=group_key,
                slide_config=slide_config,
                max_retries=args.max_retries,
                dry_run=args.dry_run,
            )
            all_results.append(slide_result)
            slide_index += 1

    total_elapsed = time.time() - total_start

    # 결과 저장
    if not args.dry_run:
        save_results(all_results, total_elapsed)

    # 최종 요약 출력
    print("\n" + "=" * 70)
    print("생성 결과 요약")
    print("=" * 70)
    success_count = sum(1 for r in all_results if r.get("success") and not args.dry_run)
    fail_count = len(all_results) - success_count if not args.dry_run else 0
    print(f"총 슬라이드: {len(all_results)}")
    if not args.dry_run:
        print(f"성공: {success_count} / 실패: {fail_count}")
        print(f"총 소요 시간: {total_elapsed:.1f}초")
        print(f"결과 파일: {RESULTS_JSON}")

    for r in all_results:
        if args.dry_run:
            continue
        status = "OK" if r.get("success") else "FAIL"
        group = r.get("group", "")
        slide = r.get("slide", "?")
        name = r.get("name", "")
        attempts = len(r.get("attempts", []))
        size = r.get("final_file_size")
        elapsed = r.get("total_elapsed_seconds")
        size_str = f"{size:,} bytes" if size else "N/A"
        time_str = f"{elapsed:.1f}초" if elapsed else "N/A"
        error_str = f" | {r['error']}" if r.get("error") else ""
        print(
            f"  [{status}] {group}/slide_{slide}_{name} | "
            f"시도: {attempts}회 | {size_str} | {time_str}{error_str}"
        )

    if args.dry_run:
        print("\n[DRY-RUN 완료] 위 프롬프트로 이미지가 생성될 예정입니다.")
        print("실제 생성하려면 --dry-run 옵션 없이 실행하세요.")


if __name__ == "__main__":
    main()