"""씬 이미지 파이프라인 — T2I 생성 + GPT LVM 검증 + I2I 연출 개선.

D. 씬 이미지 생성:
  1) Gemini T2I 생성 (참조 이미지 포함)
  2) GPT LVM 검증 → 심각도
  3) 심각하면 재생성 → GPT LVM 비교 선택
  4) 베이스 이미지 확정

E. I2I 연출 개선:
  1) GPT LVM: 구도/색감 개선안 n개 추천
  2) Gemini I2I: n개 변형 생성
  3) GPT LVM: 최종 대표 이미지 선택
"""

import json
import logging
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from app.core.config import settings
from app.modules.pipeline.ref_image_pipeline import _load_lvm_prompt
from app.modules.llm.gemini_image_client import GeminiImageClient, ModerationError
from app.modules.pipeline.ref_image_pipeline import _call_gpt_lvm

logger = logging.getLogger(__name__)

# ── GPT LVM 스키마 ──

_SCENE_VALIDATION_SCHEMA = {
    "type": "object",
    "additionalProperties": False,
    "properties": {
        "matches_prompt": {"type": "boolean"},
        "severity": {"type": "string", "description": "ok | minor | severe"},
        "issues": {"type": "array", "items": {"type": "string"}},
        "score": {"type": "integer"},
    },
    "required": ["matches_prompt", "severity", "issues", "score"],
}

_COMPARISON_SCHEMA = {
    "type": "object",
    "additionalProperties": False,
    "properties": {
        "winner": {"type": "string", "description": "image_1 | image_2"},
        "reason": {"type": "string"},
    },
    "required": ["winner", "reason"],
}

_IMPROVEMENT_SCHEMA = {
    "type": "object",
    "additionalProperties": False,
    "properties": {
        "improvements": {
            "type": "array",
            "items": {
                "type": "object",
                "additionalProperties": False,
                "properties": {
                    "type": {"type": "string", "description": "angle | color | angle+color"},
                    "prompt": {"type": "string", "description": "I2I용 텍스트 프롬프트"},
                    "reason": {"type": "string"},
                },
                "required": ["type", "prompt", "reason"],
            },
        },
    },
    "required": ["improvements"],
}

_FINAL_SELECTION_SCHEMA = {
    "type": "object",
    "additionalProperties": False,
    "properties": {
        "selected_index": {"type": "integer", "description": "0-based index"},
        "reason": {"type": "string"},
    },
    "required": ["selected_index", "reason"],
}


def generate_and_validate_scene(
    gemini_client: GeminiImageClient,
    t2i_prompt: str,
    beat_title: str,
    output_dir: Path,
    reference_images: Optional[List[Tuple[str, bytes]]] = None,
    previous_scene_bytes: Optional[bytes] = None,
) -> Dict[str, Any]:
    """씬 이미지 생성 + GPT LVM 검증 + 필요 시 재생성.

    Returns:
        {"file_path": str, "image_bytes": bytes, "validation": {...}, "was_regenerated": bool}
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    # 참조 이미지 구성
    labeled_refs = list(reference_images or [])
    if previous_scene_bytes:
        labeled_refs.append(("Previous scene (same location)", previous_scene_bytes))

    # 1) T2I 생성 — ModerationError 시 sanitizer로 프롬프트 수정 후 재시도 (최대 3회)
    logger.info("Generating scene image: %s", beat_title[:40])
    from app.modules.prompt_sanitizer import PromptSanitizer
    from app.modules.llm.openai_client import OpenAIClient

    current_prompt = t2i_prompt
    img_1 = None
    for attempt in range(4):
        try:
            img_1, _ = gemini_client.generate_image(
                prompt=current_prompt,
                labeled_references=labeled_refs if labeled_refs else None,
                aspect_ratio="16:9",
            )
            break
        except ModerationError as exc:
            logger.warning("Scene T2I blocked (attempt %d): %s", attempt + 1, exc.block_reason)
            if attempt >= 3:
                raise
            try:
                sanitizer = PromptSanitizer(OpenAIClient())
                sanitize_result = sanitizer.sanitize(current_prompt, exc.block_reason, exc.block_categories, attempt=attempt+1)
                current_prompt = sanitize_result.get("sanitized_prompt", current_prompt)
                logger.info("Sanitized scene prompt (strategy: %s)", sanitize_result.get("strategy", ""))
            except Exception:
                raise exc

    if img_1 is None:
        raise RuntimeError(f"Failed to generate scene image: {beat_title}")

    path_1 = output_dir / f"{uuid.uuid4()}.png"
    path_1.write_bytes(img_1)

    # 2) GPT LVM 검증
    try:
        validation = _call_gpt_lvm(
            img_1,
            _load_lvm_prompt("scene_validation", t2i_prompt=t2i_prompt, beat_title=beat_title),
            _SCENE_VALIDATION_SCHEMA,
            "scene_validation",
        )
    except Exception as exc:
        logger.warning("Scene LVM validation failed: %s", exc)
        validation = {"matches_prompt": True, "severity": "ok", "issues": [], "score": 70}

    # 3) severe면 재생성 + 비교
    was_regenerated = False
    if validation.get("severity") == "severe":
        logger.info("Severe — regenerating scene: %s", beat_title[:40])
        try:
            img_2, _ = gemini_client.generate_image(
                prompt=t2i_prompt,
                labeled_references=labeled_refs if labeled_refs else None,
                aspect_ratio="16:9",
            )
            path_2 = output_dir / f"{uuid.uuid4()}.png"
            path_2.write_bytes(img_2)

            comparison = _call_gpt_lvm(
                img_1,
                f"두 씬 이미지를 비교. 프롬프트에 더 맞는 쪽 선택.\n"
                f"프롬프트: {t2i_prompt}\n"
                f"image_1=첫번째, image_2=두번째",
                _COMPARISON_SCHEMA,
                "scene_comparison",
                image_bytes_2=img_2,
            )
            if comparison.get("winner") == "image_2":
                img_1 = img_2
                path_1 = path_2
                was_regenerated = True
        except Exception as exc:
            logger.warning("Scene regeneration failed: %s", exc)

    return {
        "file_path": str(path_1),
        "image_bytes": img_1,
        "validation": validation,
        "was_regenerated": was_regenerated,
        "generation_model": settings.gemini_image_model,
    }


def recommend_improvements(
    image_bytes: bytes,
    t2i_prompt: str,
    beat_title: str,
    n: int = 3,
) -> List[Dict[str, Any]]:
    """GPT LVM에게 카메라 구도/색감 개선안 추천 받기.

    Returns: [{"type": "angle"|"color"|"angle+color", "prompt": str, "reason": str}]
    """
    prompt = _load_lvm_prompt("scene_improvement", n=n, t2i_prompt=t2i_prompt, beat_title=beat_title)
    try:
        result = _call_gpt_lvm(
            image_bytes, prompt, _IMPROVEMENT_SCHEMA, "scene_improvements",
        )
        return result.get("improvements", [])[:n]
    except Exception as exc:
        logger.warning("Improvement recommendation failed: %s", exc)
        return []


def generate_i2i_variants(
    gemini_client: GeminiImageClient,
    original_bytes: bytes,
    improvements: List[Dict[str, Any]],
    output_dir: Path,
) -> List[Dict[str, Any]]:
    """개선안 기반 I2I 변형 생성.

    Returns: [{"file_path": str, "image_bytes": bytes, "improvement": {...}}]
    """
    from app.modules.gemini_i2i_editor import GeminiI2IEditor

    editor = GeminiI2IEditor(
        api_key=settings.gemini_api_key,
        model=settings.gemini_image_model,
    )
    results = []

    for imp in improvements:
        try:
            i2i_prompt = imp.get("prompt", "")
            imp_type = imp.get("type", "color")

            if imp_type == "color":
                edited = editor.edit_color(original_bytes, i2i_prompt)
            elif imp_type == "angle":
                edited = editor.edit_color(original_bytes, f"Camera angle adjustment: {i2i_prompt}")
            else:
                # angle+color — 구조화 angle params가 스키마에 없으므로 프롬프트 기반 편집
                edited = editor.edit_color(original_bytes, i2i_prompt)

            path = output_dir / f"{uuid.uuid4()}.png"
            path.write_bytes(edited)
            results.append({
                "file_path": str(path),
                "image_bytes": edited,
                "improvement": imp,
            })
        except Exception as exc:
            logger.warning("I2I variant failed: %s", exc)

    return results


def select_best_image(
    original_bytes: bytes,
    variant_images: List[bytes],
    beat_title: str,
) -> int:
    """GPT LVM이 원본 + 변형 중 최고를 선택.

    Returns: 0-based index (0=original, 1=variant[0], ...)
    """
    if not variant_images:
        return 0

    # 원본 + 변형 모두 텍스트로 설명
    descriptions = ["image_0 (original)"]
    for i in range(len(variant_images)):
        descriptions.append(f"image_{i + 1} (variant {i + 1})")

    prompt = (
        f"아래 이미지들 중 씬 '{beat_title}'의 대표 이미지로 가장 적합한 것을 선택하세요.\n\n"
        f"이미지 목록: {', '.join(descriptions)}\n"
        f"첫 번째 이미지가 원본, 나머지는 변형입니다.\n"
        f"가장 영화적이고 장면의 핵심을 잘 전달하는 이미지를 고르세요.\n"
    )

    # GPT LVM에 원본 + 첫 번째 변형만 비교 (API 제한상 이미지 2장)
    try:
        if len(variant_images) == 1:
            result = _call_gpt_lvm(
                original_bytes, prompt, _COMPARISON_SCHEMA, "final_select",
                image_bytes_2=variant_images[0],
            )
            return 0 if result.get("winner") == "image_1" else 1
        else:
            # 다수 변형: 토너먼트 방식 (원본 vs 각 변형)
            best_bytes = original_bytes
            best_idx = 0
            for i, var_bytes in enumerate(variant_images):
                result = _call_gpt_lvm(
                    best_bytes,
                    f"두 이미지 중 씬 '{beat_title}'에 더 적합한 것은?\nimage_1=현재최선, image_2=후보",
                    _COMPARISON_SCHEMA,
                    "tournament_select",
                    image_bytes_2=var_bytes,
                )
                if result.get("winner") == "image_2":
                    best_bytes = var_bytes
                    best_idx = i + 1
            return best_idx
    except Exception as exc:
        logger.warning("Final selection failed: %s", exc)
        return 0