"""Gemini 이미지 생성 클라이언트 — urllib 기반 구현 + 키 라운드로빈 + 자동 로깅."""

import base64
import json
import threading
import time
import socket
import urllib.error
import urllib.request
from typing import Dict, List, Optional, Tuple

GEMINI_API_URL_TEMPLATE = (
    "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
)
from app.core.config import settings as _settings
from app.core.image_call_budget import (
    ImageCallBudgetExceeded,
    reserve_current_call,
)
from app.modules.llm.gemini_key_pool import get_next_key
from app.modules.llm.image_format import ensure_png_bytes
from app.modules.llm.llm_logger import log_llm_call
from app.services.image_capture.sink import capture_generated_image


class ModerationError(Exception):
    """Raised when Gemini blocks content due to safety/moderation policy."""

    def __init__(self, block_reason: str, block_categories: list, raw_response: dict):
        self.block_reason = block_reason
        self.block_categories = block_categories
        self.raw_response = raw_response
        super().__init__(f"Content moderation blocked: {block_reason}")


class GeminiImageClient:
    """Gemini image generation client using the REST API + 키 라운드로빈."""

    def __init__(self, api_key: str = None, model: str = None) -> None:
        self._fixed_api_key = api_key  # 명시 전달 시 고정, None이면 라운드로빈
        self._model = model or _settings.gemini_image_model
        self._tls = threading.local()  # thread-safe context

    def set_context(self, **kwargs) -> "GeminiImageClient":
        """로깅 컨텍스트 설정 (thread-local — 동시 호출 안전)."""
        if not hasattr(self._tls, "ctx"):
            self._tls.ctx = {}
        self._tls.ctx.update(kwargs)
        return self

    @property
    def _ctx(self) -> Dict:
        """현재 스레드의 context."""
        return getattr(self._tls, "ctx", {})

    def _trace_step(self) -> str:
        """스텝 이름 — ctx > capture scope(ambient) > 'image_gen'.

        ctx 는 부르는 쪽이 채울 때만 있다. 안 채우면 전부 'image_gen' 으로
        뭉개져, 같은 실행에서 gpt 경로는 `scene_image_pipeline` 으로 남는데
        gemini 경로만 스텝을 못 갈랐다(2026-08-08 Opik 실측: 롤 86건이
        전부 image_gen). 어느 스텝이 살아 있는지가 기록의 존재 이유다.
        """
        step = self._ctx.get("step") or self._ctx.get("step_name")
        if step:
            return str(step)
        try:
            from app.modules.llm.image_tracer import (
                ambient_call_meta,
                resolve_step_name,
            )

            return resolve_step_name("image_gen", ambient_call_meta())
        except Exception:  # noqa: BLE001 — 기록 보조가 본 작업을 막지 않는다
            return "image_gen"

    # Phase 4 iter 7 I1 — fan-out 단위 trace key. 기존 column 매핑 외 모든 추가
    # 정보는 metadata dict 로 묶어 log_llm_call(metadata=...) 로 forward.
    _LOG_COLUMN_KEYS = frozenset({
        "project_id", "episode_id", "operation_type", "step_name",
        "reference_image_ids",
    })
    _LOG_METADATA_KEYS = frozenset({
        "scene_index", "shot_index", "still_id", "entity_id",
        "beat_title", "entity_name", "entity_type",
        # 2026-08-13 #108 실측 수리: set_context 가 롤/fix 태그를 싣는데
        # (make_nb2_gen_fn H2) 이 필터가 걸러 metadata_json 에 안 남아,
        # exact 태그 resolve(_resolve_generation_call_id)가 원리적으로
        # miss — 최종 자산 generation_call_id 가 214/215 NULL 이었다.
        "multiroll_tag",
        # 2026-08-16 SAFETY 사다리 (Codex BLOCK-1): 후퇴 단계명이 DB·Opik
        # metadata 에 남아야 발화 흔적 SOT 가 성립한다 — allowlist 밖이면
        # set_context 로 실어도 조용히 탈락. None 은 projection 이 걸러
        # primary 의 명시 청소(safety_ladder=None)가 메타 잡음을 안 만든다.
        "safety_ladder",
    })

    @property
    def _log_ctx(self) -> Dict:
        """log_llm_call 에 전달 가능한 context. 기존 column 단위 key 는 그대로
        forward, fan-out 단위 trace key 는 `metadata` dict 로 묶어 보낸다.
        """
        raw = self._ctx
        filtered: Dict = {
            k: v for k, v in raw.items() if k in self._LOG_COLUMN_KEYS
        }
        # step → step_name 매핑 (legacy)
        if "step" in raw and "step_name" not in filtered:
            filtered["step_name"] = raw["step"]
        if "step_name" not in filtered:
            # ctx 에 스텝이 없으면 capture scope 에서 해석 — 전에는 NULL 로
            # 남아 DB 에서도 어느 스텝의 호출인지 못 갈랐다(Opik 과 동일 병).
            resolved = self._trace_step()
            if resolved != "image_gen":
                filtered["step_name"] = resolved
        meta = {
            k: v for k, v in raw.items()
            if k in self._LOG_METADATA_KEYS and v is not None
        }
        if meta:
            filtered["metadata"] = meta
        return filtered

    @property
    def _opik_meta(self) -> Dict:
        """Phase 4 iter 7 review I1 — Opik tracer extra_metadata 용 dict.
        DB log 의 metadata 와 동일한 fan-out 단위 추적 정보 (PID/EID 포함)
        를 Opik trace metadata 에도 병합 — 운영자가 incident triage 시
        DB grep 와 Opik 검색 결과가 일관되게 보이도록.
        """
        raw = self._ctx
        meta: Dict = {}
        # column 단위 key 도 Opik metadata 에는 포함 (Opik 에는 column 구분 없음).
        for k in self._LOG_COLUMN_KEYS | self._LOG_METADATA_KEYS:
            v = raw.get(k)
            if v is None:
                continue
            meta[k] = v
        return meta

    def _get_api_key(self) -> str:
        if self._fixed_api_key:
            return self._fixed_api_key
        return get_next_key()

    def generate_image(
        self,
        prompt: str,
        reference_images: Optional[List[bytes]] = None,
        aspect_ratio: str = "16:9",
        labeled_references: Optional[List[Tuple[str, bytes]]] = None,
    ) -> Tuple[bytes, int]:
        """Generate an image from a text prompt with optional reference images.

        Returns:
            Tuple of (PNG image bytes, response_time_ms).
        """
        parts: list = []

        # 참조 이미지 ID 수집 (로깅용)
        ref_image_ids = []
        if labeled_references:
            parts.append({"text": prompt})
            for i, (label, img_bytes) in enumerate(labeled_references, 1):
                parts.append({"text": label or f"Reference image {i}:"})
                parts.append({
                    "inline_data": {
                        "mime_type": "image/png",
                        "data": base64.b64encode(img_bytes).decode("ascii"),
                    }
                })
                if label:
                    ref_image_ids.append(label)
        elif reference_images:
            parts.append({"text": prompt})
            for img_bytes in reference_images:
                parts.append({
                    "inline_data": {
                        "mime_type": "image/png",
                        "data": base64.b64encode(img_bytes).decode("ascii"),
                    }
                })
        else:
            parts.append({"text": prompt})

        body: Dict = {
            "contents": [{"parts": parts}],
            "generationConfig": {
                "responseModalities": ["TEXT", "IMAGE"],
                "imageConfig": {
                    "aspectRatio": aspect_ratio,
                    "imageSize": "1K",
                },
            },
        }

        from app.modules.llm.image_tracer import get_image_tracer
        _tracer = get_image_tracer()

        def _trace_error(err_msg, ms):
            _tracer.log(
                step=self._trace_step(), model=self._model,
                prompt=prompt, ref_image_ids=ref_image_ids or [],
                status="error", error=err_msg, duration_ms=ms,
                extra_metadata=self._opik_meta,
            )

        api_key = self._get_api_key()
        url = GEMINI_API_URL_TEMPLATE.format(model=self._model, api_key=api_key)
        req = urllib.request.Request(
            url,
            data=json.dumps(body).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )

        start_time = time.monotonic()
        last_error: Exception | None = None
        timeout = _settings.llm_timeout_image_gen
        max_retries = _settings.llm_max_retries
        for attempt in range(1, max_retries + 2):
            try:
                reserve_current_call(source="gemini_image_client.urlopen")
                with urllib.request.urlopen(req, timeout=timeout) as response:
                    payload = json.loads(response.read().decode("utf-8"))
                break
            except ImageCallBudgetExceeded:
                raise
            except urllib.error.HTTPError as exc:
                error_text = exc.read().decode("utf-8", errors="replace")
                last_error = RuntimeError(f"Gemini image API error {exc.code}: {error_text}")
                if exc.code in {429, 500, 502, 503, 504} and attempt <= max_retries:
                    # 429/503은 다른 키로 재시도 (할당량/과부하 우회)
                    if exc.code in {429, 503} and not self._fixed_api_key:
                        api_key = self._get_api_key()
                        url = GEMINI_API_URL_TEMPLATE.format(model=self._model, api_key=api_key)
                        req = urllib.request.Request(
                            url, data=json.dumps(body).encode("utf-8"),
                            headers={"Content-Type": "application/json"}, method="POST",
                        )
                    time.sleep(2 * attempt)
                    continue
                elapsed_ms = int((time.monotonic() - start_time) * 1000)
                log_llm_call(
                    model_name=self._model, user_prompt=prompt, status="error",
                    error_message=str(last_error), duration_ms=elapsed_ms,
                    reference_image_ids=ref_image_ids, **self._log_ctx,
                )
                _trace_error(str(last_error), elapsed_ms)
                raise RuntimeError(f"Gemini image API error {exc.code}: {error_text}") from exc
            except (urllib.error.URLError, socket.timeout) as exc:
                last_error = exc
                if attempt <= max_retries:
                    # timeout 시에도 다른 키로 재시도
                    if not self._fixed_api_key:
                        api_key = self._get_api_key()
                        url = GEMINI_API_URL_TEMPLATE.format(model=self._model, api_key=api_key)
                        req = urllib.request.Request(
                            url, data=json.dumps(body).encode("utf-8"),
                            headers={"Content-Type": "application/json"}, method="POST",
                        )
                    time.sleep(2 * attempt)
                    continue
                elapsed_ms = int((time.monotonic() - start_time) * 1000)
                log_llm_call(
                    model_name=self._model, user_prompt=prompt, status="error",
                    error_message=str(last_error), duration_ms=elapsed_ms,
                    reference_image_ids=ref_image_ids, **self._log_ctx,
                )
                _trace_error(str(last_error), elapsed_ms)
                raise RuntimeError(f"Gemini image API request failed after retries: {exc}") from exc
        else:
            elapsed_ms = int((time.monotonic() - start_time) * 1000)
            log_llm_call(
                model_name=self._model, user_prompt=prompt, status="error",
                error_message=str(last_error), duration_ms=elapsed_ms,
                reference_image_ids=ref_image_ids, **self._log_ctx,
            )
            _trace_error(str(last_error), elapsed_ms)
            raise RuntimeError(f"Gemini image API request failed: {last_error}")

        elapsed_ms = int((time.monotonic() - start_time) * 1000)

        # Check for moderation blocks
        try:
            self._check_moderation_block(payload)
        except ModerationError as exc:
            log_llm_call(
                model_name=self._model, user_prompt=prompt, status="error",
                error_message=str(exc), duration_ms=elapsed_ms,
                reference_image_ids=ref_image_ids, **self._log_ctx,
            )
            _tracer.log(
                step=self._trace_step(),
                model=self._model, prompt=prompt,
                ref_image_ids=ref_image_ids or [],
                status="moderation_blocked", error=str(exc),
                duration_ms=elapsed_ms,
                extra_metadata=self._opik_meta,
            )
            raise

        # Extract image bytes — text-only response 감지 (Phase 9.2 사고 fix):
        # Gemini가 거부 시 promptFeedback.blockReason이 아닌 candidate parts에 text로
        # 거부 메시지를 담아 정상 응답 형식으로 반환 (예: "I can't generate ...").
        # 이 케이스를 ModerationError로 분류해 sanitize retry 트리거.
        candidates = payload.get("candidates", [])
        text_refusal: Optional[str] = None
        for candidate in candidates:
            content = candidate.get("content", {})
            for part in content.get("parts", []):
                inline_data = part.get("inlineData") or part.get("inline_data")
                if isinstance(inline_data, dict):
                    b64 = inline_data.get("data")
                    if b64:
                        call_id = log_llm_call(
                            model_name=self._model, user_prompt=prompt,
                            output_text="[image generated]",
                            status="success", duration_ms=elapsed_ms,
                            reference_image_ids=ref_image_ids, **self._log_ctx,
                        )
                        _tracer.log(
                            step=self._trace_step(),
                            model=self._model, prompt=prompt,
                            ref_image_ids=ref_image_ids or [],
                            output_image_id=self._ctx.get("image_id"),
                            duration_ms=elapsed_ms,
                            params={"aspect_ratio": aspect_ratio},
                            extra_metadata=self._opik_meta,
                        )
                        # ★생성 바이트를 PNG 로 정규화한다 (2026-08-06).
                        #  모델이 `inlineData.mimeType: image/jpeg` 로 돌려주는
                        #  일이 실제로 있고, 그대로 `.png` 로 저장돼 왔다
                        #  (A·C 실행 생성 롤 2,037장 전건 JPEG). 이름과 내용이
                        #  어긋나면 Anthropic 판정이 400 을 내고, i2i 체인은
                        #  매 단계 재압축된다. capture 보다 **앞에서** 바꿔야
                        #  저장본과 lineage 가 같은 바이트를 가리킨다.
                        img_bytes = ensure_png_bytes(
                            base64.b64decode(b64),
                            context=f"{self._model}/"
                                    f"{self._ctx.get('operation_type') or '?'}",
                        )
                        # Phase B: 생성 이미지 capture (scope 미배선이면 no-op → 최종
                        # reference/scene 경로는 중복 0). 폐기되던 success-path
                        # log_llm_call 반환값(call_id)을 generation_call_id 로 회수.
                        # role 은 호출자 scope(operation_type)가 결정, 없으면 neutral.
                        _capture_meta = {
                            k: self._ctx.get(k)
                            for k in (
                                "project_id", "episode_id", "operation_type",
                                "step", "scene_index", "shot_index",
                                "still_id", "entity_id", "image_id",
                            )
                            if self._ctx.get(k) is not None
                        }
                        # ★ ref_image_ids 는 사람이 읽는 reference 라벨(image_asset UUID
                        # 아님) → input_image_ids(UUID lineage SOT)에 넣으면 Phase D edge
                        # 가 오염된다(Codex BLOCKING). 표시/진단 메타로만 보존하고
                        # input_image_ids 는 None(풍부 lineage 는 Phase C scope 에서).
                        if ref_image_ids:
                            _capture_meta["reference_labels"] = ref_image_ids
                        capture_generated_image(
                            img_bytes,
                            role=self._ctx.get("operation_type") or "gemini_image",
                            prompt=prompt,
                            generation_call_id=call_id,
                            pipeline_metadata=_capture_meta,
                        )
                        return img_bytes, elapsed_ms
                # 이미지 없는 part — text 응답 보존 (거부 메시지 추출용)
                if "text" in part and part.get("text"):
                    text_refusal = part["text"]

        # No image extracted. text 응답이 있으면 거부로 간주 → ModerationError로
        # 던져 ref_image_pipeline의 sanitize retry 루프(최대 3회)가 작동.
        if text_refusal:
            finish_reason = (
                candidates[0].get("finishReason", "STOP") if candidates else "UNKNOWN"
            )
            log_llm_call(
                model_name=self._model, user_prompt=prompt, status="error",
                error_message=f"text_refusal: {text_refusal[:200]}",
                duration_ms=elapsed_ms,
                reference_image_ids=ref_image_ids, **self._log_ctx,
            )
            _tracer.log(
                step=self._trace_step(),
                model=self._model, prompt=prompt,
                ref_image_ids=ref_image_ids or [],
                status="moderation_blocked_text",
                error=text_refusal[:200], duration_ms=elapsed_ms,
                extra_metadata=self._opik_meta,
            )
            raise ModerationError(
                block_reason=f"text_refusal:{finish_reason}",
                block_categories=["TEXT_REFUSAL"],
                raw_response=payload,
            )

        # 본문이 아예 비어 있고 **사유만** 오는 거절 (2026-08-18 실측):
        #   {"candidates":[{"content":{},"finishReason":"PROHIBITED_CONTENT",...}]}
        # 위 text_refusal 분기는 거부 문구가 part 로 올 때만 걸리므로 이 모양은
        # 그동안 일반 RuntimeError 로 나갔다. 그러면 부르는 쪽이 검열로 못
        # 알아보고 순화 재시도·후퇴 사다리가 통째로 건너뛰어져 샷이 죽는다
        # (S5sh2 실측: 사다리가 교차 백엔드에서 멈췄다). 사유가 거절일 때만
        # ModerationError 로 올린다 — 그 밖의 빈 응답은 기존대로 일반 오류.
        # ★IMAGE_PROHIBITED_CONTENT 는 이번 결함과 같은 모양의 이미지 정책
        #  거절이다(litellm 의 Gemini finish-reason 지도도 content_filter 로
        #  선언한다 — litellm_core_utils/core_helpers.py:90). 빠뜨리면 같은
        #  구멍이 그대로 남는다. STOP·MAX_TOKENS·OTHER 는 일부러 뺀다 —
        #  거절이 아닌 것까지 넣으면 유료 후퇴가 헛돈다.
        _refusal_reasons = {
            "PROHIBITED_CONTENT", "IMAGE_PROHIBITED_CONTENT",
            "IMAGE_SAFETY", "SAFETY", "RECITATION", "BLOCKLIST", "SPII",
        }
        _finish = (candidates[0].get("finishReason") or "") if candidates else ""
        if str(_finish).upper() in _refusal_reasons:
            log_llm_call(
                model_name=self._model, user_prompt=prompt, status="error",
                error_message=f"finish_refusal: {_finish}",
                duration_ms=elapsed_ms,
                reference_image_ids=ref_image_ids, **self._log_ctx,
            )
            _tracer.log(
                step=self._trace_step(),
                model=self._model, prompt=prompt,
                ref_image_ids=ref_image_ids or [],
                status="moderation_blocked_finish",
                error=str(_finish)[:200], duration_ms=elapsed_ms,
                extra_metadata=self._opik_meta,
            )
            raise ModerationError(
                block_reason=f"finish_refusal:{_finish}",
                block_categories=[str(_finish)],
                raw_response=payload,
            )

        _trace_error("No image parts in response", elapsed_ms)
        log_llm_call(
            model_name=self._model, user_prompt=prompt, status="error",
            error_message="No image parts in response", duration_ms=elapsed_ms,
            reference_image_ids=ref_image_ids, **self._log_ctx,
        )
        raise RuntimeError(f"Gemini image API returned no image parts: {json.dumps(payload)[:500]}")

    @staticmethod
    def _check_moderation_block(payload: dict) -> None:
        """Detect moderation blocks in the Gemini response."""
        prompt_feedback = payload.get("promptFeedback", {})
        block_reason = prompt_feedback.get("blockReason")
        if block_reason:
            categories = [
                r.get("category", "UNKNOWN")
                for r in prompt_feedback.get("safetyRatings", [])
                if r.get("probability", "NEGLIGIBLE") not in ("NEGLIGIBLE", "LOW")
            ]
            raise ModerationError(
                block_reason=block_reason,
                block_categories=categories,
                raw_response=payload,
            )

        candidates = payload.get("candidates", [])
        if candidates:
            finish_reason = candidates[0].get("finishReason", "")
            if finish_reason in ("SAFETY", "IMAGE_SAFETY", "IMAGE_OTHER"):
                safety_ratings = candidates[0].get("safetyRatings", [])
                categories = [
                    r.get("category", "UNKNOWN")
                    for r in safety_ratings
                    if r.get("probability", "NEGLIGIBLE") not in ("NEGLIGIBLE", "LOW")
                ]
                raise ModerationError(
                    block_reason="SAFETY",
                    block_categories=categories,
                    raw_response=payload,
                )

        if not candidates:
            safety_ratings = prompt_feedback.get("safetyRatings", [])
            high_risk = [
                r.get("category", "UNKNOWN")
                for r in safety_ratings
                if r.get("probability", "NEGLIGIBLE") not in ("NEGLIGIBLE", "LOW")
            ]
            if high_risk:
                raise ModerationError(
                    block_reason="SAFETY",
                    block_categories=high_risk,
                    raw_response=payload,
                )