"""
Step 18: Contextual Retrieval — 메뉴별 식별 문맥 생성 및 embedding_text 보강

[Anthropic Contextual Retrieval, 2024]
  각 메뉴에 대해 GPT가 "같은 카테고리의 유사 메뉴들과 구별되는 1~2문장 식별 문맥"을 생성.
  이 문맥을 embedding_text 앞에 prepend 후 벡터 재임베딩.
  → 비슷한 메뉴 간 혼동 감소, 검색 정확도 향상 (Anthropic 발표: ~47% 실패 감소)

실행:
    python scripts/18_contextual_retrieval.py

완료 후 실행:
    python scripts/02_build_vectordb.py --reset
    python scripts/06_rebuild_bm25.py
    python scripts/17_eval_comparison.py
"""
import asyncio
import json
import shutil
import subprocess
import sys
from datetime import datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from openai import AsyncAzureOpenAI
from config import AZURE_KEY, AZURE_ENDPOINT, AZURE_API_VERSION, LLM_MODEL

# ── 경로 ──────────────────────────────────────────────────────────────────────
BASE_DIR   = Path(__file__).parent.parent
JSONL_PATH = BASE_DIR / "data" / "generated" / "menu_descriptions.jsonl"
BACKUP_PATH = JSONL_PATH.with_name("menu_descriptions.jsonl.pre_contextual")

# ── 파라미터 ──────────────────────────────────────────────────────────────────
CONCURRENCY   = 15      # 동시 API 호출 (이미지 없음 → 높게 설정)
MAX_SIBLINGS  = 25      # 형제 메뉴 최대 표시 개수
CONTEXT_MARKER = "【식별문맥】"  # 재실행 시 기존 문맥 제거를 위한 마커


# ── 유틸 ──────────────────────────────────────────────────────────────────────

def load_menus() -> list[dict]:
    with open(JSONL_PATH, encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]


def save_menus(menus: list[dict]) -> None:
    with open(JSONL_PATH, "w", encoding="utf-8") as f:
        for m in menus:
            f.write(json.dumps(m, ensure_ascii=False) + "\n")


def strip_existing_context(text: str) -> str:
    """이미 추가된 식별 문맥 제거 (멱등성 보장)"""
    if CONTEXT_MARKER in text:
        idx = text.find("\n\n", text.find(CONTEXT_MARKER))
        if idx != -1:
            return text[idx + 2:]
    return text


def build_sibling_map(menus: list[dict]) -> dict[str, list[str]]:
    """2-depth 카테고리 키 → 메뉴명 목록"""
    result: dict[str, list[str]] = {}
    for m in menus:
        parts = [p.strip() for p in m.get("menu_path", "").split(">")]
        key = " > ".join(parts[:2]) if len(parts) >= 2 else (parts[0] if parts else "기타")
        result.setdefault(key, []).append(m["menu_name"])
    return result


# ── 비동기 GPT 호출 ──────────────────────────────────────────────────────────

CONTEXT_SYSTEM = (
    "당신은 증권 모바일 앱(영웅문S#) 메뉴 검색 시스템 전문가입니다. "
    "주어진 메뉴가 비슷한 메뉴들과 어떻게 구별되는지 명확하게 서술하세요."
)


def _build_context_prompt(menu: dict, sibling_names: list[str]) -> str:
    others = [n for n in sibling_names if n != menu["menu_name"]][:MAX_SIBLINGS]
    sibling_str = ", ".join(others) if others else "없음"

    return (
        f"같은 카테고리의 다른 메뉴들: {sibling_str}\n\n"
        f"메뉴 경로: {menu['menu_path']}\n"
        f"메뉴명: {menu['menu_name']}\n"
        f"기능 설명: {menu.get('function_desc', '')}\n"
        f"키워드: {', '.join(menu.get('keywords', [])[:15])}\n\n"
        f"위 메뉴에 대한 식별 문맥을 1~2문장으로 작성하세요.\n"
        f"[조건]\n"
        f"- 이 메뉴만의 핵심 기능과 사용 상황을 명확히 서술할 것\n"
        f"- 동일 카테고리의 비슷한 메뉴(예: 수익 관련 4개 메뉴)와 무엇이 다른지 구체적으로 언급할 것\n"
        f"- ~습니다 / ~입니다 말투로 끝낼 것\n"
        f"- JSON 없이 텍스트만 출력할 것"
    )


async def _gen_one(
    client: AsyncAzureOpenAI,
    semaphore: asyncio.Semaphore,
    menu: dict,
    sibling_names: list[str],
    idx: int,
    total: int,
) -> str:
    prompt = _build_context_prompt(menu, sibling_names)
    async with semaphore:
        try:
            resp = await client.chat.completions.create(
                model=LLM_MODEL,
                messages=[
                    {"role": "system", "content": CONTEXT_SYSTEM},
                    {"role": "user",   "content": prompt},
                ],
                max_tokens=160,
                temperature=0.1,
            )
            ctx = resp.choices[0].message.content.strip()
            print(f"  [{idx+1:3d}/{total}] {menu['menu_name'][:18]:<18} → {ctx[:50]}…")
            return ctx
        except Exception as e:
            print(f"  [{idx+1:3d}/{total}] ⚠ {menu['menu_name']}: {e}")
            return ""


async def _run_all(menus: list[dict], sibling_map: dict[str, list[str]]) -> list[str]:
    client = AsyncAzureOpenAI(
        api_key=AZURE_KEY,
        azure_endpoint=AZURE_ENDPOINT,
        api_version=AZURE_API_VERSION,
    )
    sem = asyncio.Semaphore(CONCURRENCY)
    tasks = []
    for i, menu in enumerate(menus):
        parts = [p.strip() for p in menu.get("menu_path", "").split(">")]
        key = " > ".join(parts[:2]) if len(parts) >= 2 else (parts[0] if parts else "기타")
        tasks.append(_gen_one(client, sem, menu, sibling_map.get(key, []), i, len(menus)))
    results = await asyncio.gather(*tasks)
    await client.close()
    return list(results)


# ── 메인 ─────────────────────────────────────────────────────────────────────

def main():
    print(f"[18_ctx] Contextual Retrieval 시작 - {datetime.now():%Y-%m-%d %H:%M:%S}")
    print(f"[18_ctx] JSONL: {JSONL_PATH}")

    # ── 1. 로드 ──────────────────────────────────────────────────────────────
    menus = load_menus()
    print(f"[18_ctx] {len(menus)}개 메뉴 로드")

    # ── 2. 백업 ──────────────────────────────────────────────────────────────
    shutil.copy2(JSONL_PATH, BACKUP_PATH)
    print(f"[18_ctx] 백업 완료: {BACKUP_PATH.name}")

    # ── 3. 기존 식별 문맥 제거 (재실행 멱등성) ───────────────────────────────
    already = sum(1 for m in menus if CONTEXT_MARKER in m.get("embedding_text", ""))
    if already:
        print(f"[18_ctx] 기존 식별 문맥 {already}개 제거 중…")
    for m in menus:
        m["embedding_text"] = strip_existing_context(m.get("embedding_text", ""))

    # ── 4. 형제 메뉴 맵 ──────────────────────────────────────────────────────
    sibling_map = build_sibling_map(menus)
    print(f"[18_ctx] 카테고리 그룹: {len(sibling_map)}개")

    # ── 5. GPT 식별 문맥 생성 (비동기 병렬) ──────────────────────────────────
    print(f"\n[18_ctx] GPT 식별 문맥 생성 (concurrency={CONCURRENCY}) …")
    contexts = asyncio.run(_run_all(menus, sibling_map))

    # ── 6. embedding_text 보강 ───────────────────────────────────────────────
    enriched = 0
    for menu, ctx in zip(menus, contexts):
        if ctx:
            menu["embedding_text"] = (
                f"{CONTEXT_MARKER} {ctx}\n\n{menu['embedding_text']}"
            )
            enriched += 1
    print(f"\n[18_ctx] 식별 문맥 추가: {enriched}/{len(menus)}개")

    # ── 7. 저장 ──────────────────────────────────────────────────────────────
    save_menus(menus)
    print(f"[18_ctx] 저장 완료: {JSONL_PATH}")

    # ── 8. 재빌드 안내 + 자동 실행 ───────────────────────────────────────────
    print("\n" + "=" * 65)
    print(" ChromaDB & BM25 재빌드 시작...")
    print("=" * 65)

    py = sys.executable
    steps = [
        [py, str(BASE_DIR / "scripts" / "02_build_vectordb.py"), "--reset"],
        [py, str(BASE_DIR / "scripts" / "06_rebuild_bm25.py")],
    ]
    for cmd in steps:
        label = Path(cmd[1]).name
        print(f"\n▶ {label} 실행 중…")
        result = subprocess.run(cmd, capture_output=False, text=True)
        if result.returncode != 0:
            print(f"  ⚠ {label} 실패 (returncode={result.returncode})")
            print("  수동으로 실행하세요:")
            print(f"    {' '.join(cmd)}")
        else:
            print(f"  ✓ {label} 완료")

    print("\n" + "=" * 65)
    print(" Contextual Retrieval 완료!")
    print("=" * 65)
    print(f"  식별 문맥 추가: {enriched}개 / {len(menus)}개")
    print(f"  백업:          {BACKUP_PATH.name}")
    print()
    print("  평가 비교 실행:")
    print("    python scripts/17_eval_comparison.py")
    print("=" * 65)


if __name__ == "__main__":
    main()