"""요소 참조 이미지 파이프라인 — T2I 생성 + GPT LVM 검증 + 비교 선택.

흐름:
1) 의존성 순서대로 T2I 생성 (Gemini) — 변형은 기본 참조이미지 포함
2) GPT LVM 검증 → 심각도 판단
3) 심각하면 재생성 → GPT LVM이 (1)과 (2) 비교하여 더 좋은 것 선택
4) 최종 참조 이미지 확정
"""

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

from app.core.config import settings
from app.modules.prompt_loader import load_prompt
from app.modules.llm.gemini_image_client import GeminiImageClient, ModerationError

# entity_type → prompt template filename (without .md)
_REF_PROMPT_MAP = {
    "character": "character_ref",
    "character_nonhuman": "character_nonhuman_ref",
    "location": "location_ref",
    "prop": "prop_ref",
    "outlook": "character_outlook_ref",
    "composite": "character_composite_ref",
}

# entity_type → aspect ratio
_ASPECT_MAP = {
    "character": "3:4",
    "character_nonhuman": "1:1",
    "outlook": "3:4",
    "prop": "1:1",
    "location": "16:9",
    "composite": "16:9",
}


def _load_ref_image_prompt(entity_type: str, **kwargs) -> Optional[str]:
    """Load type-specific reference image prompt template from external files.

    Returns the formatted prompt string, or None if no template is found
    (backward-compatible: caller falls back to raw t2i_prompt).
    """
    template_name = _REF_PROMPT_MAP.get(entity_type)
    if not template_name:
        return None
    try:
        return load_prompt("ref_image_prompts", template_name, **kwargs)
    except (FileNotFoundError, KeyError):
        return None


def _load_lvm_prompt(name: str, **kwargs) -> str:
    return load_prompt("lvm_prompts", name, **kwargs)

logger = logging.getLogger(__name__)

OPENAI_API_URL = "https://api.openai.com/v1/responses"

# ── GPT LVM 스키마 ──

_VALIDATION_SCHEMA = {
    "type": "object",
    "additionalProperties": False,
    "properties": {
        "matches_description": {"type": "boolean"},
        "severity": {
            "type": "string",
            "description": "ok | minor | severe",
        },
        "issues": {
            "type": "array",
            "items": {"type": "string"},
        },
        "score": {"type": "integer", "description": "0-100"},
    },
    "required": ["matches_description", "severity", "issues", "score"],
}

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


def _call_gpt_lvm(
    image_bytes: bytes,
    text_prompt: str,
    response_schema: Dict[str, Any],
    schema_name: str = "lvm_result",
    image_bytes_2: Optional[bytes] = None,
    opik_tags: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """GPT LVM (Vision) 호출 — LiteLLM Router 경유, Opik 자동 추적."""
    from app.modules.llm.llm_client import router_completion

    content = [
        {"type": "text", "text": text_prompt},
        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64.b64encode(image_bytes).decode('ascii')}"}},
    ]
    if image_bytes_2:
        content.append(
            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{base64.b64encode(image_bytes_2).decode('ascii')}"}}
        )

    tags = opik_tags if opik_tags else [schema_name]

    # Phase 4 iter 7 W1 — temperature param 명시 전달 제거. gpt-5*/gpt-5.4*
    # family 는 temperature=1 만 허용하고 0.2 같은 값은 OpenAI 측에서 직접
    # 400 BadRequest 로 거부함. router/litellm 의 drop_params 가 response_format
    # path 에서 일관되게 적용되지 않아 silent BadRequest → silent absorb 사고
    # 가 발생했음 (운영 결함, scene/ref 모든 호출). 호출 측에서 보내지 않는
    # 것이 안전 — vision validation 의 의미상 default 가 적합.
    response = router_completion(
        model="gpt",
        messages=[{"role": "user", "content": content}],
        response_format={
            "type": "json_schema",
            "json_schema": {"name": schema_name, "schema": response_schema, "strict": True},
        },
        timeout=settings.llm_timeout_validation,
        metadata={"opik": {"tags": tags}},
    )

    raw = response.choices[0].message.content
    if not raw:
        raise RuntimeError(f"GPT LVM returned empty response for schema={schema_name}")
    return json.loads(raw)


def validate_reference_image(
    image_bytes: bytes,
    entity_name: str,
    entity_description: str,
    entity_type: str,
) -> Dict[str, Any]:
    """GPT LVM으로 참조 이미지 검증.

    Returns: {"matches_description": bool, "severity": "ok"|"minor"|"severe", "issues": [...], "score": 0-100}
    """
    prompt = _load_lvm_prompt("ref_validation",
        entity_name=entity_name, entity_type=entity_type, entity_description=entity_description)
    return _call_gpt_lvm(image_bytes, prompt, _VALIDATION_SCHEMA, "ref_validation",
                         opik_tags=["ref_validation"])


def compare_two_images(
    image_1: bytes,
    image_2: bytes,
    entity_name: str,
    entity_description: str,
) -> Dict[str, Any]:
    """GPT LVM으로 두 이미지 비교하여 더 나은 것 선택.

    Returns: {"winner": "image_1"|"image_2", "reason": "..."}
    """
    prompt = _load_lvm_prompt("ref_comparison",
        entity_name=entity_name, entity_description=entity_description)
    return _call_gpt_lvm(image_1, prompt, _COMPARISON_SCHEMA, "ref_comparison",
                         image_bytes_2=image_2, opik_tags=["ref_comparison"])


def generate_and_validate_reference(
    gemini_client: GeminiImageClient,
    entity_name: str,
    entity_description: str,
    entity_type: str,
    t2i_prompt: str,
    output_dir: Path,
    extra_references: Optional[List[Tuple[str, bytes]]] = None,
    style_context: str = "",
    trace_meta: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """참조 이미지 생성 + GPT LVM 검증 + 필요 시 재생성 비교.

    Args:
        trace_meta (Phase 4 iter 7 W3): observability metadata —
            {project_id, episode_id, entity_id, scene_index, shot_index}.
            llm_call_log 의 PID/EID NULL 사고 차단용.

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

    # Opik context — W3: trace_meta 를 set_context 에 forward.
    ctx_kwargs: Dict[str, Any] = {
        "step": "ref_image_gen",
        "entity_name": entity_name,
        "entity_type": entity_type,
    }
    if trace_meta:
        ctx_kwargs.update(trace_meta)
    gemini_client.set_context(**ctx_kwargs)

    # 1) T2I 생성 — ModerationError 시 sanitizer로 프롬프트 수정 후 재시도 (최대 3회)
    logger.info("Generating reference image for: %s (type=%s)", entity_name, entity_type)
    from app.modules.prompt_sanitizer import PromptSanitizer
    from app.modules.llm.openai_client import OpenAIClient

    # 타입별 프롬프트 템플릿 로드 (외부 파일 우선, 없으면 raw t2i_prompt 사용)
    ref_prompt = _load_ref_image_prompt(
        entity_type,
        entity_description=entity_description,
        outlook_description=entity_description,
    )
    if ref_prompt:
        current_prompt = ref_prompt
        logger.info("Using type-specific ref prompt template for %s (%s)", entity_name, entity_type)
    else:
        current_prompt = t2i_prompt

    # 스타일 컨텍스트를 프롬프트 앞에 추가
    if style_context:
        current_prompt = f"{style_context}\n\n{current_prompt}"
    aspect = _ASPECT_MAP.get(entity_type, "3:4")
    img_bytes_1 = None
    sanitization_info = None

    for attempt in range(4):  # 1 original + 3 sanitized retries
        try:
            img_bytes_1, elapsed = gemini_client.generate_image(
                prompt=current_prompt,
                labeled_references=extra_references,
                aspect_ratio=aspect,
            )
            break
        except ModerationError as exc:
            logger.warning("Reference T2I blocked (attempt %d) for %s: %s",
                           attempt + 1, entity_name, exc.block_reason)
            if attempt >= 3:
                raise
            # Sanitize prompt and retry
            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)
                sanitization_info = sanitize_result
                logger.info("Sanitized prompt for %s (strategy: %s): %s",
                            entity_name, sanitize_result.get("strategy", ""), current_prompt[:80])
            except Exception as san_exc:
                logger.warning("Sanitization failed: %s", san_exc)
                raise exc

    if img_bytes_1 is None:
        raise RuntimeError(f"Failed to generate reference image for {entity_name}")

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

    # PNG 메타데이터 삽입
    from app.modules.png_metadata import embed_png_metadata
    from datetime import datetime, timezone

    embed_png_metadata(path_1, {
        "prompt": current_prompt,
        "entity_name": entity_name,
        "entity_type": entity_type,
        "model": settings.gemini_image_model,
        "created_at": datetime.now(timezone.utc).isoformat(),
    })

    # 2) GPT LVM 검증
    # Phase 4 iter 7 W1 — silent OK fallback 제거 (scene_image_pipeline 와 동일).
    # default AppError fail-fast — silent absorb 차단. 운영자 override env =
    # ALLOW_LVM_VALIDATION_UNAVAILABLE=true.
    try:
        validation = validate_reference_image(
            img_bytes_1, entity_name, entity_description, entity_type,
        )
    except Exception as exc:
        if settings.allow_lvm_validation_unavailable:
            logger.warning(
                "Reference LVM validation unavailable for %s: %s — explicit "
                "override (ALLOW_LVM_VALIDATION_UNAVAILABLE=true). Proceeding "
                "with _validation_unavailable marker.",
                entity_name, exc,
            )
            # Phase 4 iter 7 I2 — matches_description=False 로 strong negative
            # signal. 옛 True 값은 silent 성공 오판 위험.
            validation = {
                "matches_description": False,
                "severity": "unavailable",
                "issues": ["LVM validation unavailable — explicit override active"],
                "score": 0,
                "_validation_unavailable": True,
            }
        else:
            from app.core.errors import AppError
            raise AppError(
                code="step.ref_image_pipeline.lvm_validation_failed",
                message=(
                    f"Reference LVM validation failed for {entity_name!r}: "
                    f"{exc!r}. To proceed without LVM (e.g. transient outage) "
                    f"set ALLOW_LVM_VALIDATION_UNAVAILABLE=true env. Silent "
                    f"absorb 금지 (feedback_no_silent_fallback)."
                ),
                status_code=500,
            )

    # 3) 심각하면 재생성 + 비교
    if validation.get("severity") == "severe":
        logger.info("Severe issue for %s — regenerating", entity_name)
        try:
            img_bytes_2, _ = gemini_client.generate_image(
                prompt=current_prompt,
                labeled_references=extra_references,
                aspect_ratio=aspect,
            )
            path_2 = output_dir / f"{uuid.uuid4()}.png"
            path_2.write_bytes(img_bytes_2)

            # GPT LVM 비교
            comparison = compare_two_images(
                img_bytes_1, img_bytes_2, entity_name, entity_description,
            )
            winner = comparison.get("winner", "image_1")
            logger.info("Comparison for %s: winner=%s reason=%s",
                        entity_name, winner, comparison.get("reason", ""))

            if winner == "image_2":
                return {
                    "file_path": str(path_2),
                    "image_bytes": img_bytes_2,
                    "prompt_used": current_prompt,
                    "validation": validation,
                    "comparison": comparison,
                    "was_regenerated": True,
                    "generation_model": settings.gemini_image_model,
                    "sanitization_info": sanitization_info,
                }
        except Exception as exc:
            logger.warning("Regeneration failed for %s: %s", entity_name, exc)

    return {
        "file_path": str(path_1),
        "image_bytes": img_bytes_1,
        "prompt_used": current_prompt,
        "validation": validation,
        "was_regenerated": False,
        "generation_model": settings.gemini_image_model,
        "sanitization_info": sanitization_info,
    }