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

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

import base64
import json
import logging
import socket
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from app.core.config import settings

_PROMPTS_BASE = Path(__file__).resolve().parent.parent.parent.parent.parent / "prompts" / "_base"

LVM_PROMPT_DIR = _PROMPTS_BASE / "lvm_prompts"

REF_IMAGE_PROMPT_DIR = _PROMPTS_BASE / "ref_image_prompts"

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


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).
    """
    if not REF_IMAGE_PROMPT_DIR.exists():
        return None
    versions = sorted(
        [d.name for d in REF_IMAGE_PROMPT_DIR.iterdir() if d.is_dir()],
        reverse=True,
    )
    if not versions:
        return None
    template_name = _REF_PROMPT_MAP.get(entity_type)
    if not template_name:
        return None
    template_path = REF_IMAGE_PROMPT_DIR / versions[0] / f"{template_name}.md"
    if not template_path.exists():
        return None
    text = template_path.read_text(encoding="utf-8").strip()
    try:
        return text.format(**kwargs) if kwargs else text
    except KeyError:
        return text

def _load_lvm_prompt(name: str, **kwargs) -> str:
    versions = sorted([d.name for d in LVM_PROMPT_DIR.iterdir() if d.is_dir()], reverse=True)
    if not versions:
        raise FileNotFoundError(f"No LVM prompt versions in {LVM_PROMPT_DIR}")
    text = (LVM_PROMPT_DIR / versions[0] / f"{name}.md").read_text(encoding="utf-8").strip()
    return text.format(**kwargs) if kwargs else text
from app.modules.llm.gemini_image_client import GeminiImageClient, ModerationError

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,
) -> Dict[str, Any]:
    """GPT LVM (Vision) 호출 — 이미지 1~2장 + 텍스트 프롬프트."""
    api_key = settings.openai_api_key
    model = settings.openai_model

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

    body = {
        "model": model,
        "input": [
            {"type": "message", "role": "user", "content": content_parts},
        ],
        "text": {
            "format": {
                "type": "json_schema",
                "name": schema_name,
                "strict": True,
                "schema": response_schema,
            }
        },
        "temperature": 0.2,
        "store": False,
    }

    req = urllib.request.Request(
        OPENAI_API_URL,
        data=json.dumps(body).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    timeout = settings.llm_timeout_validation
    max_retries = settings.llm_max_retries
    last_error = None

    for attempt in range(1, max_retries + 2):
        try:
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                payload = json.loads(resp.read().decode("utf-8"))
            break
        except urllib.error.HTTPError as exc:
            error_text = exc.read().decode("utf-8", errors="replace")
            last_error = RuntimeError(f"GPT LVM error {exc.code}: {error_text}")
            if exc.code in {429, 500, 502, 503, 504} and attempt <= max_retries:
                time.sleep(2 * attempt)
                continue
            raise last_error from exc
        except (urllib.error.URLError, socket.timeout) as exc:
            last_error = exc
            if attempt <= max_retries:
                time.sleep(2 * attempt)
                continue
            raise RuntimeError(f"GPT LVM failed: {exc}") from exc
    else:
        raise RuntimeError(f"GPT LVM failed: {last_error}")

    # Extract result
    output_text = payload.get("output_text", "")
    if not output_text:
        for item in payload.get("output", []):
            if isinstance(item, dict):
                for part in item.get("content", []):
                    if isinstance(part, dict) and part.get("type") == "output_text":
                        output_text = part.get("text", "")
                        break
    return json.loads(output_text)


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")


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)


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 = "",
) -> Dict[str, Any]:
    """참조 이미지 생성 + GPT LVM 검증 + 필요 시 재생성 비교.

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

    # 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 = "1:1" if entity_type == "character" else "16:9"
    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)

    # 2) GPT LVM 검증
    try:
        validation = validate_reference_image(
            img_bytes_1, entity_name, entity_description, entity_type,
        )
    except Exception as exc:
        logger.warning("LVM validation failed for %s: %s", entity_name, exc)
        validation = {"matches_description": True, "severity": "ok", "issues": [], "score": 70}

    # 3) 심각하면 재생성 + 비교
    if validation.get("severity") == "severe":
        logger.info("Severe issue for %s — regenerating", entity_name)
        try:
            img_bytes_2, _ = gemini_client.generate_image(
                prompt=t2i_prompt,
                labeled_references=extra_references,
                aspect_ratio="1:1" if entity_type == "character" else "16:9",
            )
            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,
                    "validation": validation,
                    "comparison": comparison,
                    "was_regenerated": True,
                    "generation_model": settings.gemini_image_model,
                }
        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,
        "validation": validation,
        "was_regenerated": False,
        "generation_model": settings.gemini_image_model,
    }