"""shot_essence_extraction StepRunner — shot description을 essence/peripheral/atmospheric 3분류.

Phase 1b (v3 plan 2026-04-29). scene_detail이 모든 시각요소를 한 t2i_prompt에
넣어 카탈로그식 합성을 일으키는 문제를 해결하기 위해, shot description을 사전
분류하여 scene_detail은 essence만 prepend (peripheral/atmospheric은 다른
step이 처리).

설계:
- 입력: shot_validator + shot_selection — selected shot 만 처리
- 출력: per shot {essence, peripheral, atmospheric}
- 처리: shots를 BUNDLE_MAX (15)로 묶어 LLM 호출, MAX_WORKERS (4) 병렬
- 모델: gpt-5.5
- toggle: settings.shot_essence_enabled — applicability="if_shot_essence_enabled"
  (default False, run-all에서 자동 제외)

Phase 2 통합 지점:
- backend/app/core/steps/scene_context_loader.py — essence 체크포인트 로드
- backend/app/core/steps/detail_steps.py::SceneDetailStep._analyze_one —
  user_prompt 앞에 essence prepend (atmospheric은 chain bg ref가 처리)

체크포인트 출력:
  data.shots[].{scene_index, shot_index, essence: [str], peripheral: [str],
                atmospheric: [str], status: "ok"|"failed"}
  data.applicable_count: 입력 shot 수 (총 처리 대상)
  data.succeeded_count:  성공 shot 수 (status='ok')
  data.failed_count:     실패 shot 수 (status='failed')
"""
import json
import logging
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from app.core.step_runner import StepRunner
from app.modules.llm.llm_client import call_structured
from app.modules.prompt_loader import load_prompt, load_schema

logger = logging.getLogger(__name__)

BUNDLE_MAX_SHOTS = 15
MAX_WORKERS = 4
MAX_RETRY = 3
RETRY_BACKOFF_BASE = 2  # exponential: 2, 4, 6 sec


class ShotEssenceExtractionStep(StepRunner):
    """Shot별 essence/peripheral/atmospheric 3분류 추출."""

    def _load_prev_checkpoint(self, step_id: str) -> Optional[Dict]:
        from app.core.config import settings
        cp = (
            Path(settings.projects_dir) / self.project_id
            / "checkpoints" / "episodes" / self.episode_id
            / step_id / "manifest.json"
        )
        if cp.exists():
            return json.loads(cp.read_text(encoding="utf-8"))
        return None

    def _execute(self, mode="resume") -> Dict[str, Any]:
        # shot_validator + shot_selection 로드
        shot_cp = self._load_prev_checkpoint("shot_validator")
        if not shot_cp or not shot_cp.get("data", {}).get("scenes"):
            from app.core.errors import AppError
            raise AppError(
                code="step.no_shots",
                message="shot_validator 결과 없음",
                status_code=400,
            )

        sel_cp = self._load_prev_checkpoint("shot_selection")
        selected_map: Dict[int, set] = {}
        if sel_cp and sel_cp.get("data", {}).get("scenes"):
            for s in sel_cp["data"]["scenes"]:
                selected_map[s["scene_index"]] = set(s.get("selected_shot_indices", []))

        # selected shots만 추출. scene_consistency 패턴: shot_selection에 누락된 씬은
        # 명시적 "선택 없음" 으로 처리 (전체 처리 X — Codex H1 회귀 가드).
        all_shots: List[Dict[str, Any]] = []
        for s in shot_cp["data"]["scenes"]:
            si = s["scene_index"]
            if si not in selected_map:
                logger.warning(
                    "shot_essence_extraction: scene %d missing from shot_selection — skipped",
                    si,
                )
                continue
            sel = selected_map[si]
            for sh in s.get("shots", []):
                if sh.get("shot_index") in sel:
                    all_shots.append({
                        "scene_index": si,
                        "shot_index": sh.get("shot_index", 0),
                        "description": sh.get("description", "") or "",
                    })

        if not all_shots:
            return {
                "completed_count": 0,
                "applicable_count": 0,
                "failed_count": 0,
                "data": {
                    "shots": [],
                    "applicable_count": 0,
                    "succeeded_count": 0,
                    "failed_count": 0,
                },
            }

        system_prompt = load_prompt("shot_essence_extraction", "system", db=self.db)
        schema = load_schema("shot_essence_extraction", "schema", db=self.db)

        # 번들 구성 (15 shots씩)
        bundles: List[List[Dict[str, Any]]] = []
        for i in range(0, len(all_shots), BUNDLE_MAX_SHOTS):
            bundles.append(all_shots[i:i + BUNDLE_MAX_SHOTS])

        logger.info(
            "shot_essence_extraction: %d selected shots → %d bundles",
            len(all_shots), len(bundles),
        )

        results_by_idx: Dict[int, List[Dict[str, Any]]] = {}
        with ThreadPoolExecutor(max_workers=MAX_WORKERS) as pool:
            futures = {
                pool.submit(self._call_bundle, bundle, i + 1, system_prompt, schema): i + 1
                for i, bundle in enumerate(bundles)
            }
            for fut in as_completed(futures):
                call_idx, shots = fut.result()
                results_by_idx[call_idx] = shots

        # call_idx 순서로 결과 합치기
        all_results: List[Dict[str, Any]] = []
        for i in sorted(results_by_idx):
            all_results.extend(results_by_idx[i])

        # ── 카운팅 ──
        # _call_bundle이 stub(failed)도 추가하므로 all_results 길이 == 입력 shot 수.
        succeeded = sum(1 for r in all_results if r.get("status") == "ok")
        failed = sum(1 for r in all_results if r.get("status") == "failed")

        if failed:
            failed_keys = [
                f"S{r['scene_index']}_Shot{r['shot_index']}"
                for r in all_results if r.get("status") == "failed"
            ]
            logger.warning(
                "shot_essence_extraction: %d shots failed: %s",
                failed, ", ".join(failed_keys),
            )

        logger.info(
            "shot_essence_extraction: %d/%d shots succeeded (%d failed)",
            succeeded, len(all_shots), failed,
        )

        return {
            "completed_count": succeeded,
            "applicable_count": len(all_shots),
            "failed_count": failed,
            "data": {
                "shots": all_results,
                "applicable_count": len(all_shots),
                "succeeded_count": succeeded,
                "failed_count": failed,
            },
        }

    def _call_bundle(
        self,
        bundle: List[Dict[str, Any]],
        call_idx: int,
        system_prompt: str,
        schema: dict,
    ) -> Tuple[int, List[Dict[str, Any]]]:
        """단일 번들 LLM 호출. retry 포함. 최종 실패 시 stub(status='failed') 반환.

        반환 결과는 항상 입력 bundle과 동일한 (scene_index, shot_index) set.
        LLM 응답에 누락/extra/중복이 있어도 입력 keys 기준으로 정렬 + stub 채움.
        """
        input_keys = {(sh["scene_index"], sh["shot_index"]) for sh in bundle}
        shots_lines: List[str] = []
        for sh in bundle:
            shots_lines.append(
                f"\n--- Scene {sh['scene_index']} Shot {sh['shot_index']} ---\n"
                f"description: {sh['description']}"
            )
        user_prompt = (
            "아래 샷들의 description을 essence/peripheral/atmospheric 3분류로 나누세요.\n"
            "각 샷마다 scene_index/shot_index/essence/peripheral/atmospheric 모두 출력.\n"
            "\n[분석 대상 샷]"
            + "\n".join(shots_lines)
        )

        scene_range = f"S{bundle[0]['scene_index']}-S{bundle[-1]['scene_index']}"
        last_exc: Optional[Exception] = None
        for attempt in range(MAX_RETRY + 1):
            try:
                result = call_structured(
                    step="shot_essence_extraction",
                    system_prompt=system_prompt,
                    user_prompt=user_prompt,
                    response_schema=schema,
                    project_config=self.project_config,
                    schema_name=f"shot_essence_{call_idx}",
                    opik_metadata=self.build_opik_metadata(),
                )
                raw_shots = result.get("shots", [])
                merged = self._merge_with_input_keys(bundle, input_keys, raw_shots)
                logger.info(
                    "shot_essence_extraction call %d (%s, %d shots): %d ok / %d failed",
                    call_idx, scene_range, len(bundle),
                    sum(1 for r in merged if r.get("status") == "ok"),
                    sum(1 for r in merged if r.get("status") == "failed"),
                )
                return call_idx, merged
            except Exception as exc:
                last_exc = exc
                if attempt < MAX_RETRY:
                    delay = RETRY_BACKOFF_BASE * (attempt + 1)
                    logger.warning(
                        "shot_essence_extraction call %d (%s) retry %d/%d (sleep %ds): %s",
                        call_idx, scene_range, attempt + 1, MAX_RETRY, delay, exc,
                    )
                    time.sleep(delay)
                    continue
                break

        # final fail — stub 채워 모든 입력 shot이 결과에 등장하도록
        logger.warning(
            "shot_essence_extraction call %d (%s) FAILED after %d retries: %s",
            call_idx, scene_range, MAX_RETRY, last_exc,
        )
        stubs = [
            {
                "scene_index": sh["scene_index"],
                "shot_index": sh["shot_index"],
                "essence": [],
                "peripheral": [],
                "atmospheric": [],
                "status": "failed",
            }
            for sh in bundle
        ]
        return call_idx, stubs

    def _merge_with_input_keys(
        self,
        bundle: List[Dict[str, Any]],
        input_keys: set,
        raw_shots: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """LLM 응답을 입력 bundle 키 set과 1:1 매칭. 누락/extra/중복 모두 처리.

        - 입력에 없는 키: 무시 + warning
        - 응답 중복: 첫 번째만 사용 + warning
        - 응답 누락: stub(status='failed') 채움
        """
        by_key: Dict[Tuple[int, int], Dict[str, Any]] = {}
        for r in raw_shots:
            si = r.get("scene_index")
            shi = r.get("shot_index")
            if si is None or shi is None:
                logger.warning("shot_essence_extraction: response item missing scene/shot index — drop")
                continue
            key = (int(si), int(shi))
            if key not in input_keys:
                logger.warning(
                    "shot_essence_extraction: extra key in response (S%d Shot%d) — drop",
                    key[0], key[1],
                )
                continue
            if key in by_key:
                logger.warning(
                    "shot_essence_extraction: duplicate response for S%d Shot%d — keep first",
                    key[0], key[1],
                )
                continue
            by_key[key] = {
                "scene_index": key[0],
                "shot_index": key[1],
                "essence": list(r.get("essence", [])),
                "peripheral": list(r.get("peripheral", [])),
                "atmospheric": list(r.get("atmospheric", [])),
                "status": "ok",
            }

        # 입력 순서대로 결과 빌드 + 누락된 key는 failed stub
        merged: List[Dict[str, Any]] = []
        for sh in bundle:
            key = (sh["scene_index"], sh["shot_index"])
            if key in by_key:
                merged.append(by_key[key])
            else:
                logger.warning(
                    "shot_essence_extraction: response missing S%d Shot%d — stub failed",
                    key[0], key[1],
                )
                merged.append({
                    "scene_index": key[0],
                    "shot_index": key[1],
                    "essence": [],
                    "peripheral": [],
                    "atmospheric": [],
                    "status": "failed",
                })
        return merged