"""인물 추출 비교 테스트 — Gemini 3.1 Pro Preview vs GPT 5.4

씬 구분된 시나리오 전문을 두 모델에 동일한 프롬프트로 보내서
추출된 인물 목록 + 출현 씬 수를 비교한다.
"""
import json
import sys
import os
from concurrent.futures import ThreadPoolExecutor, as_completed

# project root를 path에 추가
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"))

# Gemini key → GOOGLE_API_KEY (litellm expects this)
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 = "시나리오 분석 전문가. 시나리오를 읽고 등장 인물을 추출한다."

USER_PROMPT_TEMPLATE = """아래 시나리오를 읽고, 등장하는 인물을 모두 나열하세요.

## 규칙
- 되도록 인물의 이름을 추가해야 함
- 엑스트라 제외 (행인, 군중, 이름 없는 단역 등 누구든 대체 가능한 인물)
- 인물이 바뀌어도 시각적 일관성에 문제 없는 인물 제외
- 각 인물이 몇 개의 씬에 등장하는지 세세요

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

SCHEMA = {
    "type": "object",
    "properties": {
        "characters": {
            "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": ["characters"],
    "additionalProperties": False,
}

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


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


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

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

    # 현재 파이프라인 결과 로드
    PID = "1d036235-39f3-419f-8267-3df305ebdfaf"
    EID = "70d7f162-2486-4838-a9ff-3c673c427722"
    BASE = f"/Users/manta/Documents/Projects/TheRoad-I1/projects/{PID}/checkpoints/episodes/{EID}"
    with open(f"{BASE}/entity_extract_character/manifest.json") as f:
        current = json.load(f)["data"]["characters"]
    current_names = {c["name"] for c in current}

    # 두 모델 병렬 호출
    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] = {"characters": []}

    # 결과 비교 출력
    print()
    print("=" * 80)
    print(f"{'인물':15s} | {'Gemini Pro':>12s} | {'GPT 5.4':>12s} | {'현재 파이프라인':>14s}")
    print("-" * 80)

    # 모든 이름 합치기
    all_names = set()
    gemini_map = {}
    gpt_map = {}

    for c in results.get("gemini", {}).get("characters", []):
        gemini_map[c["name"]] = c
        all_names.add(c["name"])

    for c in results.get("gpt", {}).get("characters", []):
        gpt_map[c["name"]] = c
        all_names.add(c["name"])

    all_names |= current_names

    for name in sorted(all_names):
        g = gemini_map.get(name)
        p = gpt_map.get(name)
        cur = "O" if name in current_names else "-"

        g_str = f"{g['scene_count']}씬" if g else "-"
        p_str = f"{p['scene_count']}씬" if p else "-"

        print(f"  {name:15s} | {g_str:>12s} | {p_str:>12s} | {cur:>14s}")

    print("-" * 80)
    print(f"  {'합계':15s} | {len(gemini_map):>11d}명 | {len(gpt_map):>11d}명 | {len(current_names):>13d}명")


if __name__ == "__main__":
    main()