"""소품 추출 비교 테스트 — Gemini 3.1 Pro Preview vs GPT 5.4"""
import json
import sys
import os
from concurrent.futures import ThreadPoolExecutor, as_completed

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from dotenv import load_dotenv
load_dotenv(os.path.join(os.path.dirname(__file__), "..", ".env"))
if os.environ.get("GEMINI_API_KEY") and not os.environ.get("GOOGLE_API_KEY"):
    os.environ["GOOGLE_API_KEY"] = os.environ["GEMINI_API_KEY"]

from litellm import completion

SCENARIO_PATH = "/tmp/yokai_scenes.md"

SYSTEM_PROMPT = "시나리오 분석 전문가. 시나리오를 읽고 이미지 생성(T2I)에 필요한 시각적 요소만 추출한다."

USER_PROMPT_TEMPLATE = """아래 시나리오를 읽고, 시각적 일관성을 유지해야 하는 중요 물체(소품)를 추출하세요.

## 포함 기준
- 화면에 중요하게 나타나는 것 (1개 씬이라도 포함)
- 작은 물체라도 사람이 컷을 볼때에 쉽게 주목되어 볼 수 있는 것

## 분할 기준
- 큰 물체는 내부와 외부, 내부의 특정 공간 외부의 특정 부분등은 다른 물체로 취급하고 외부가 아닌경우 상세 부분 표기필요(내부, 특정 부분 등으로)

## 합치기 규칙
- 같은 기능의 변형은 하나로

## 제외 기준
- 프롬프트만으로 여러 번 물체 이미지를 생성했을 때 차이점을 사람이 쉽게 구분하기 어렵거나 못 느낄 정도면 제외
- 사람이 입는 형태의 모든 것 (수트, 우주복, 갑옷, 제복, 의상, 입는 장치나 로봇은 제외)
- 배경에 설치된 장비 (벽면 모니터, TV, CCTV 등)
- 배경/건물의 일부 (문, 창문, 계단, 엘리베이터)
- UI/화면 표시 (상태창, 모니터 화면, HUD)

## 목표
- 에피소드당 5~10개 (절대 15개 초과 금지)
- 각 소품이 몇 개의 씬에 등장하는지, 어떤 씬인지 세세요

## 시나리오 전문
{scenario}
"""

SCHEMA = {
    "type": "object",
    "properties": {
        "props": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "name": {"type": "string"},
                    "scene_count": {"type": "integer"},
                    "scene_indices": {
                        "type": "array",
                        "items": {"type": "integer"},
                    },
                },
                "required": ["name", "scene_count", "scene_indices"],
                "additionalProperties": False,
            },
        }
    },
    "required": ["props"],
    "additionalProperties": False,
}

MODELS = {
    "gemini": "gemini/gemini-3.1-pro-preview",
    "gpt": "gpt-5.5",
}


def call_model(label, model, scenario):
    resp = completion(
        model=model,
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": USER_PROMPT_TEMPLATE.format(scenario=scenario)},
        ],
        response_format={
            "type": "json_schema",
            "json_schema": {"name": "props", "schema": SCHEMA, "strict": True},
        },
        temperature=0.2,
    )
    return json.loads(resp.choices[0].message.content)


def main():
    with open(SCENARIO_PATH, "r") as f:
        scenario = f.read()

    print(f"시나리오: {len(scenario):,} chars\n")

    # 기존 프로젝트 결과 로드
    import glob
    BASE = "/Users/manta/Documents/Projects/TheRoad-I1/projects"
    yokai = {"7b0117e5": "요괴전 v9", "1d036235": "요괴전 v13"}
    for pid_prefix, name in yokai.items():
        matches = glob.glob(f"{BASE}/{pid_prefix}*")
        if not matches:
            continue
        cp = glob.glob(f"{matches[0]}/checkpoints/episodes/*/entity_extract_prop/manifest.json")
        if not cp:
            print(f"[{name}] 소품 체크포인트 없음")
            continue
        with open(cp[0]) as f:
            data = json.load(f)
        props = data.get("data", {}).get("props", [])
        print(f"[{name}] 기존 {len(props)}개: {[p['name'] for p in props]}")
    print()

    results = {}
    with ThreadPoolExecutor(max_workers=2) as pool:
        futures = {
            pool.submit(call_model, label, model, scenario): label
            for label, model in MODELS.items()
        }
        for fut in as_completed(futures):
            label = futures[fut]
            try:
                results[label] = fut.result()
                print(f"[{label}] 완료")
            except Exception as e:
                print(f"[{label}] 실패: {e}")
                results[label] = {"props": []}

    gemini_props = {p["name"]: p for p in results.get("gemini", {}).get("props", [])}
    gpt_props = {p["name"]: p for p in results.get("gpt", {}).get("props", [])}

    all_names = sorted(set(gemini_props.keys()) | set(gpt_props.keys()))

    print()
    print(f"{'소품':25s} | {'Gemini Pro':>12s} | {'GPT 5.4':>12s}")
    print("-" * 60)

    for name in all_names:
        g = gemini_props.get(name)
        p = gpt_props.get(name)
        g_str = f"{g['scene_count']}씬" if g else "-"
        p_str = f"{p['scene_count']}씬" if p else "-"
        print(f"  {name:25s} | {g_str:>12s} | {p_str:>12s}")

    print("-" * 60)
    print(f"  {'합계':25s} | {len(gemini_props):>11d}개 | {len(gpt_props):>11d}개")


if __name__ == "__main__":
    main()