"""배경 추출 비교 테스트 — 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 = "시나리오 분석 전문가. 시나리오를 읽고 등장하는 배경(장소)을 추출한다."

USER_PROMPT_TEMPLATE = """아래 시나리오를 읽고, 등장하는 모든 배경(장소)을 나열하세요.

## 분리/통합 기준 — 영화 세트 기준
- 같은 건물/장소의 안과 밖은 분리 (시각적으로 다른 세트)
- 같은 내부 공간(방, 복도, 거실 등)은 하나로 통합 (같은 세트)
- 지붕, 꼭대기 등 시각적으로 완전히 다른 공간은 분리
- 영화에서 같은 세트에서 촬영할 수준이면 하나로 묶어라

## 규칙
- 1개 씬에만 등장해도 포함
- 각 배경이 몇 개의 씬에 등장하는지, 어떤 씬인지 세세요

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

SCHEMA = {
    "type": "object",
    "properties": {
        "locations": {
            "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": ["locations"],
    "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": "locations", "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")

    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] = {"locations": []}

    gemini_locs = {l["name"]: l for l in results.get("gemini", {}).get("locations", [])}
    gpt_locs = {l["name"]: l for l in results.get("gpt", {}).get("locations", [])}

    all_names = sorted(set(gemini_locs.keys()) | set(gpt_locs.keys()))

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

    for name in all_names:
        g = gemini_locs.get(name)
        p = gpt_locs.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_locs):>11d}개 | {len(gpt_locs):>11d}개")


if __name__ == "__main__":
    main()