"""Phase 9.1 비교 실험 — primary shot 이미지를 다른 모델로 재생성.

의도: production scene_image_pipeline은 Gemini Image (`gemini-3.1-flash-image-preview`)
사용. 이 스크립트는 동일한 ref + t2i_prompt를 다른 모델 (gpt-image-2 등)에
보내 결과를 비교한다.

특징:
- production 코드 (scene_image_pipeline) 안 건드림 — 일회성
- 모델 client 추상화 (`ImageGenClient` 인터페이스) — 미래 swap 쉽게
- 결과 PNG는 별도 디렉토리 (`scene_<model_label>_phase91/`)
- DB 업데이트 X (비교만 — DB는 Gemini 결과 보존)

사용:
    cd backend && .venv/bin/python ../scripts/regen_phase91_with_model.py \\
        --pid c00bbe19-a9b5-463f-acfc-806f2e820258 \\
        --ep fe165e3a-19c2-4a0f-9acb-e0c9bab0ee5a \\
        --model gpt-image-2 \\
        --out-label gpt_image_2_phase91
"""
from __future__ import annotations

import argparse
import base64
import io
import json
import logging
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import List, Optional, Protocol, Tuple

PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT / "backend"))

from app.core.database import SessionLocal
from app.models.project import ImageAsset, SceneStill
from sqlalchemy import and_

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s %(message)s",
)
logger = logging.getLogger("regen_phase91")


# ---------------------------------------------------------------------------
# Image Generation Client Abstraction (model swap)
# ---------------------------------------------------------------------------

class ImageGenClient(Protocol):
    """간단한 ref + prompt → PNG bytes 인터페이스. 모델별 구현체가 따른다."""

    label: str

    def generate(
        self, *, prompt: str, ref_image_paths: List[Path], size: str = "1024x1024"
    ) -> bytes:
        """PNG bytes 반환. 실패 시 RuntimeError."""
        ...


class GeminiImageGenClient:
    """Gemini Image (`gemini-3.1-flash-image-preview` = nano banana 2.0)."""

    label = "gemini"

    def __init__(self, aspect_ratio: str = "16:9"):
        from app.modules.llm.gemini_image_client import GeminiImageClient
        self._client = GeminiImageClient()
        self._aspect = aspect_ratio

    def generate(
        self, *, prompt: str, ref_image_paths: List[Path], size: str = "1024x1024"
    ) -> bytes:
        ref_bytes_list: List[Tuple[str, bytes]] = []
        for i, p in enumerate(ref_image_paths, 1):
            if p.exists():
                ref_bytes_list.append((f"Reference image {i}", p.read_bytes()))
        png_bytes, _ms = self._client.generate_image(
            prompt=prompt,
            labeled_references=ref_bytes_list or None,
            aspect_ratio=self._aspect,
        )
        return png_bytes


class OpenAIImageGenClient:
    """OpenAI gpt-image-2 (또는 gpt-image-1) via images.edit / images.generate.

    `moderation="low"`은 GPT Image 시리즈의 콘텐츠 필터를 가장 덜 엄격하게 설정.
    auto(default) / low 두 값만 지원.

    Phase 9.2: 항상 landscape (16:9 근접) — caller의 size 인자 무시하고 1536x1024
    사용. gpt-image-2/1 표준 지원 사이즈 중 가장 wide. Gemini가 default 16:9이고
    이 client는 비교 실험용이므로 production과 framing 비교 가능하게 맞춘다.
    """

    DEFAULT_LANDSCAPE_SIZE = "1536x1024"  # 3:2 — gpt-image 표준 중 16:9 최근접

    def __init__(
        self,
        model: str = "gpt-image-2",
        quality: str = "high",
        moderation: str = "low",
    ):
        self.label = f"{model}_mod_{moderation}"
        self._model = model
        self._quality = quality
        self._moderation = moderation
        try:
            from app.modules.llm.llm_client import get_openai_client
            self._client = get_openai_client()
        except Exception:
            from openai import OpenAI
            self._client = OpenAI()

    def generate(
        self, *, prompt: str, ref_image_paths: List[Path], size: str = "1024x1024"
    ) -> bytes:
        valid = [p for p in ref_image_paths if p.exists()]

        # 16:9 정책 — caller의 size 무시하고 landscape 강제. caller가 다른 size를
        # 명시하면 warning 한 번 발생 (silent override 방지).
        if size != self.DEFAULT_LANDSCAPE_SIZE:
            logger.warning(
                "OpenAIImageGenClient: caller size=%r overridden to %r (16:9 정책)",
                size, self.DEFAULT_LANDSCAPE_SIZE,
            )
        size = self.DEFAULT_LANDSCAPE_SIZE

        # openai SDK 2.28은 `moderation` kwarg 미지원 → extra_body로 raw 전달.
        common_kwargs = {
            "model": self._model,
            "prompt": prompt,
            "size": size,
            "quality": self._quality,
            "n": 1,
            "extra_body": {"moderation": self._moderation},
        }

        if len(valid) >= 2:
            import contextlib
            with contextlib.ExitStack() as stack:
                files = [stack.enter_context(p.open("rb")) for p in valid]
                resp = self._client.images.edit(image=files, **common_kwargs)
        elif len(valid) == 1:
            with valid[0].open("rb") as f:
                resp = self._client.images.edit(image=f, **common_kwargs)
        else:
            resp = self._client.images.generate(**common_kwargs)

        b64 = resp.data[0].b64_json if resp and resp.data else None
        if not b64:
            raise RuntimeError("empty b64 response")
        return base64.b64decode(b64)


def make_client(model: str) -> ImageGenClient:
    if model == "gemini":
        return GeminiImageGenClient()
    if model in ("gpt-image-2", "gpt-image-1"):
        return OpenAIImageGenClient(model=model)
    raise ValueError(f"unknown model: {model}")


# ---------------------------------------------------------------------------
# Data extraction — 각 primary shot의 prompt + ref paths
# ---------------------------------------------------------------------------

def _resolve_ref_paths(
    db, primary: ImageAsset, projects_dir: Path
) -> List[Path]:
    """primary shot ImageAsset의 reference_image_ids를 풀어 file_path list 반환."""
    ref_ids: List[str] = []
    try:
        ref_ids = json.loads(primary.reference_image_ids or "[]")
    except Exception:
        ref_ids = []

    paths: List[Path] = []
    for rid in ref_ids:
        ref_asset = db.query(ImageAsset).filter(ImageAsset.id == rid).first()
        if ref_asset and ref_asset.file_path:
            p = Path(ref_asset.file_path)
            if not p.is_absolute():
                p = projects_dir / p
            if p.exists():
                paths.append(p)
    return paths


def _build_chain_bg_map(
    *, project_id: str, episode_id: str, projects_dir: Path
) -> dict:
    """background_render cp groups[bg_id] → (si, shi) → png_path 매핑.

    Phase 9.1 scene_image_pipeline이 inject하는 chain_bg ref와 동일.
    close_skip은 호출자가 별도로 적용 (camera_direction 정규식)."""
    import re as _re
    sid_re = _re.compile(r"^S(\d+)_Shot(\d+)$")
    cp_path = (
        projects_dir / project_id / "checkpoints" / "episodes" / episode_id
        / "background_render" / "manifest.json"
    )
    if not cp_path.exists():
        return {}
    cp = json.loads(cp_path.read_text(encoding="utf-8"))
    groups = (cp.get("data") or {}).get("groups") or {}
    result: dict = {}
    for bg_id, g in groups.items():
        if g.get("status") != "ok":
            continue
        png_path = g.get("png_path")
        if not png_path:
            continue
        for sid in g.get("shot_ids") or []:
            m = sid_re.match(sid)
            if not m:
                continue
            key = (int(m.group(1)), int(m.group(2)))
            if key not in result:
                p = Path(png_path)
                if not p.is_absolute():
                    p = projects_dir / p
                if p.exists():
                    result[key] = p
    return result


_CLOSE_FRAMING_RE = None  # lazy

def _is_close_framing(camera_direction: str) -> bool:
    """scene_generation_coordinator와 동일 정규식 — close skip 결정."""
    global _CLOSE_FRAMING_RE
    if _CLOSE_FRAMING_RE is None:
        from app.services.scene_generation_coordinator import _CLOSE_FRAMING_RE as _RE
        _CLOSE_FRAMING_RE = _RE
    if not camera_direction:
        return False
    return bool(_CLOSE_FRAMING_RE.search(camera_direction))


def _build_staging_map(
    *, project_id: str, episode_id: str, projects_dir: Path
) -> dict:
    """shot_staging cp → (si, shi) → camera_direction string."""
    cp_path = (
        projects_dir / project_id / "checkpoints" / "episodes" / episode_id
        / "shot_staging" / "manifest.json"
    )
    if not cp_path.exists():
        return {}
    cp = json.loads(cp_path.read_text(encoding="utf-8"))
    shots = (cp.get("data") or {}).get("shots") or []
    result = {}
    for sh in shots:
        si = sh.get("scene_index"); shi = sh.get("shot_index")
        if si is None or shi is None:
            continue
        result[(si, shi)] = sh.get("camera_direction") or ""
    return result


def collect_shot_jobs(
    *, project_id: str, episode_id: str, projects_dir: Path
) -> List[dict]:
    """primary shot ImageAsset + 그 prompt + entity refs + chain_bg ref(close_skip 적용)."""
    chain_bg_map = _build_chain_bg_map(
        project_id=project_id, episode_id=episode_id, projects_dir=projects_dir
    )
    staging_map = _build_staging_map(
        project_id=project_id, episode_id=episode_id, projects_dir=projects_dir
    )
    db = SessionLocal()
    try:
        rows = (
            db.query(ImageAsset)
            .filter(and_(
                ImageAsset.project_id == project_id,
                ImageAsset.episode_id == episode_id,
                ImageAsset.asset_type == "scene",
                ImageAsset.is_primary == 1,
            ))
            .all()
        )
        jobs: List[dict] = []
        chain_bg_used = 0
        chain_bg_skipped_close = 0
        for asset in rows:
            still = (
                db.query(SceneStill)
                .filter(SceneStill.id == asset.still_id)
                .first()
            )
            if not still:
                continue
            si, shi = still.scene_index, still.shot_index
            entity_refs = _resolve_ref_paths(db, asset, projects_dir)
            ref_paths: List[Path] = list(entity_refs)
            prompt = asset.prompt_used or ""
            if not prompt:
                continue
            # chain_bg ref injection — Phase 9.1 close_skip 적용
            cam_dir = staging_map.get((si, shi), "")
            chain_bg_path = chain_bg_map.get((si, shi))
            ref_kind = "entity_only"
            if chain_bg_path and not _is_close_framing(cam_dir):
                ref_paths.insert(0, chain_bg_path)
                chain_bg_used += 1
                ref_kind = "chain_bg"
            elif chain_bg_path and _is_close_framing(cam_dir):
                chain_bg_skipped_close += 1
                ref_kind = "close_skip"

            jobs.append({
                "still_id": still.id,
                "scene_index": si,
                "shot_index": shi,
                "shot_description": still.shot_description or "",
                "prompt": prompt,
                "ref_paths": ref_paths,
                "ref_kind": ref_kind,
                "original_path": asset.file_path,
            })
        jobs.sort(key=lambda j: (j["scene_index"] or 0, j["shot_index"] or 0))
        logger.info(
            "ref injection: chain_bg=%d, close_skip=%d, total=%d",
            chain_bg_used, chain_bg_skipped_close, len(jobs),
        )
        return jobs
    finally:
        db.close()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--pid", required=True)
    parser.add_argument("--ep", required=True)
    parser.add_argument("--model", required=True,
                        help="gemini | gpt-image-2 | gpt-image-1")
    parser.add_argument("--out-label", required=True,
                        help="출력 디렉토리 suffix (예: gpt_image_2_phase91)")
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--projects-dir", default="projects")
    args = parser.parse_args()

    projects_dir = (PROJECT_ROOT / args.projects_dir).resolve()
    out_dir = projects_dir / args.pid / "images" / args.ep / f"scene_{args.out_label}"
    out_dir.mkdir(parents=True, exist_ok=True)

    jobs = collect_shot_jobs(
        project_id=args.pid, episode_id=args.ep, projects_dir=projects_dir
    )
    logger.info("collected %d primary shot jobs", len(jobs))

    if args.dry_run:
        for j in jobs[:3]:
            print(f"S{j['scene_index']}_Shot{j['shot_index']}: refs={len(j['ref_paths'])} prompt_len={len(j['prompt'])}")
        return

    client = make_client(args.model)
    logger.info("using client: %s", client.label)

    completed = 0
    failed = 0
    t0 = time.time()

    def _process(job):
        out_path = out_dir / f"S{job['scene_index']:02d}_Shot{job['shot_index']:02d}.png"
        if out_path.exists():
            return ("skip", job, "already exists")
        try:
            png = client.generate(
                prompt=job["prompt"],
                ref_image_paths=job["ref_paths"],
                size="1024x1024",
            )
            out_path.write_bytes(png)
            return ("ok", job, str(out_path))
        except Exception as exc:
            return ("fail", job, str(exc)[:300])

    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {pool.submit(_process, j): j for j in jobs}
        for fut in as_completed(futures):
            status, job, info = fut.result()
            sid = f"S{job['scene_index']}_Shot{job['shot_index']}"
            if status == "ok":
                completed += 1
                logger.info("%s OK (%d/%d) -> %s", sid, completed, len(jobs), info)
            elif status == "skip":
                completed += 1
                logger.info("%s SKIP (%d/%d): %s", sid, completed, len(jobs), info)
            else:
                failed += 1
                logger.error("%s FAIL: %s", sid, info)

    elapsed = int(time.time() - t0)
    logger.info("done. completed=%d failed=%d elapsed=%ds out=%s",
                completed, failed, elapsed, out_dir)


if __name__ == "__main__":
    main()