"""
평가셋 300개 구축 스크립트
===========================
top200 메뉴에서 270개 (메뉴별 균등 분산) + 201위이하에서 30개 = 300개

핵심 원칙: 뽑힌 300개는 쿼리 인덱스에서 제거 (데이터 누수 방지)

실행:
    .venv/Scripts/python.exe scripts/21_build_eval300.py
    .venv/Scripts/python.exe scripts/21_build_eval300.py --seed 42  # 재현 가능
    .venv/Scripts/python.exe scripts/21_build_eval300.py --dry-run  # 통계만
"""
import sys
import re
import json
import random
import argparse
from pathlib import Path
from collections import defaultdict

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

from openpyxl import load_workbook
from openpyxl.cell.rich_text import CellRichText, TextBlock

EXCEL_PATH  = Path(__file__).parent.parent / "data" / "MenuSearch_review_all_20260525_edit.xlsx"
EVAL_OUT    = Path(__file__).parent.parent / "data" / "eval_queries_300.json"
INDEX_REMAINING_OUT = Path(__file__).parent.parent / "data" / "query_index_pool.json"

TOP200_N    = 270   # top200에서 뽑을 수
OUTSIDE_N   = 30    # 201+에서 뽑을 수
LLM_COL     = 8
HUMAN_COL   = 9


def parse_q(text: str) -> list[str]:
    if not text or str(text).strip() in ("", "nan"):
        return []
    text = str(text).strip()
    parts = re.split(r"\n?\s*\d+[.)]\s+", text)
    if len(parts) > 1:
        return [p.strip() for p in parts if len(p.strip()) >= 4]
    return [p.strip() for p in re.split(r"\s*/\s*|\n", text) if len(p.strip()) >= 4]


def extract_kept(val) -> list[str]:
    """취소선 제외하고 유효 텍스트만 반환"""
    if val is None:
        return []
    if isinstance(val, CellRichText):
        kept = []
        for block in val:
            if isinstance(block, TextBlock):
                if not (block.font and block.font.strike):
                    kept.append(str(block.text))
            elif isinstance(block, str):
                kept.append(block)
        return parse_q("".join(kept))
    return parse_q(str(val))


def load_all_queries(excel_path: Path) -> tuple[list[dict], list[dict]]:
    """top200 쿼리와 201+ 쿼리를 분리해서 반환"""
    wb = load_workbook(excel_path, rich_text=True)
    ws = wb.active

    top200, outside = [], []

    for row_idx in range(2, ws.max_row + 1):
        rank      = ws.cell(row_idx, 1).value
        menu_id   = ws.cell(row_idx, 2).value
        menu_name = ws.cell(row_idx, 3).value
        menu_path = ws.cell(row_idx, 4).value

        if not menu_id or str(menu_id).strip() in ("", "nan"):
            continue

        try:
            rank_int = int(rank)
        except (TypeError, ValueError):
            rank_int = 9999

        menu_id   = str(menu_id).strip()
        menu_name = str(menu_name).strip() if menu_name else ""
        menu_path = str(menu_path).strip() if menu_path else ""

        llm_qs   = extract_kept(ws.cell(row_idx, LLM_COL).value)
        human_qs = extract_kept(ws.cell(row_idx, HUMAN_COL).value)

        for q in human_qs:
            item = dict(query=q, menu_id=menu_id, menu_name=menu_name,
                        menu_path=menu_path, rank=rank_int, source="human")
            (top200 if rank_int <= 200 else outside).append(item)

        for q in llm_qs:
            item = dict(query=q, menu_id=menu_id, menu_name=menu_name,
                        menu_path=menu_path, rank=rank_int, source="llm")
            (top200 if rank_int <= 200 else outside).append(item)

    return top200, outside


def stratified_sample(queries: list[dict], n: int, rng: random.Random) -> list[dict]:
    """
    메뉴별로 균등하게 n개 샘플링.
    각 메뉴에서 최소 1개 보장 후, 남은 슬롯은 쿼리 수 비례 배분.
    human 쿼리를 llm보다 우선 선택.
    """
    # 메뉴별 그룹화 (human 먼저 정렬)
    by_menu: dict[str, list[dict]] = defaultdict(list)
    for item in queries:
        by_menu[item["menu_id"]].append(item)

    # 각 메뉴 내에서 human 우선 셔플
    for mid in by_menu:
        human = [x for x in by_menu[mid] if x["source"] == "human"]
        llm   = [x for x in by_menu[mid] if x["source"] == "llm"]
        rng.shuffle(human)
        rng.shuffle(llm)
        by_menu[mid] = human + llm  # human 먼저

    menus = list(by_menu.keys())
    n_menus = len(menus)
    n = min(n, len(queries))

    # 1단계: 메뉴별 최소 1개
    selected: list[dict] = []
    pool_by_menu: dict[str, list[dict]] = {m: list(v) for m, v in by_menu.items()}

    for mid in menus:
        if pool_by_menu[mid]:
            selected.append(pool_by_menu[mid].pop(0))
        if len(selected) >= n:
            break

    # 2단계: 남은 슬롯 채우기 (쿼리 많은 메뉴 우선)
    remaining_flat = []
    for mid in menus:
        remaining_flat.extend(pool_by_menu[mid])
    rng.shuffle(remaining_flat)

    for item in remaining_flat:
        if len(selected) >= n:
            break
        selected.append(item)

    rng.shuffle(selected)
    return selected


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--seed",    type=int, default=2025, help="랜덤 시드")
    parser.add_argument("--dry-run", action="store_true",    help="통계만 출력, 파일 저장 안 함")
    args = parser.parse_args()

    rng = random.Random(args.seed)

    print(f"엑셀 로드: {EXCEL_PATH}")
    top200_pool, outside_pool = load_all_queries(EXCEL_PATH)
    print(f"  top200 풀: {len(top200_pool):,}개 ({len(set(x['menu_id'] for x in top200_pool))}개 메뉴)")
    print(f"  201+  풀: {len(outside_pool):,}개 ({len(set(x['menu_id'] for x in outside_pool))}개 메뉴)")

    # ── 샘플링 ────────────────────────────────────────────────────────────────
    eval_top200  = stratified_sample(top200_pool, TOP200_N, rng)
    eval_outside = rng.sample(outside_pool, min(OUTSIDE_N, len(outside_pool)))

    eval_set = eval_top200 + eval_outside
    rng.shuffle(eval_set)

    # eval ID 부여
    for i, item in enumerate(eval_set, 1):
        item["id"] = i

    # ── 인덱스용 나머지 (eval 제외한 top200만) ───────────────────────────────
    eval_ids = {(x["query"], x["menu_id"]) for x in eval_set}
    index_pool = [x for x in top200_pool if (x["query"], x["menu_id"]) not in eval_ids]

    # ── 통계 ──────────────────────────────────────────────────────────────────
    eval_menus = len(set(x["menu_id"] for x in eval_set))
    top200_menus_in_eval = len(set(x["menu_id"] for x in eval_top200))
    outside_menus_in_eval = len(set(x["menu_id"] for x in eval_outside))
    human_cnt = sum(1 for x in eval_set if x["source"] == "human")
    llm_cnt   = sum(1 for x in eval_set if x["source"] == "llm")

    print(f"\n[평가셋 구성]")
    print(f"  총 쿼리:       {len(eval_set)}개")
    print(f"  top200 쿼리:   {len(eval_top200)}개  ({top200_menus_in_eval}개 메뉴)")
    print(f"  201+  쿼리:    {len(eval_outside)}개  ({outside_menus_in_eval}개 메뉴)")
    print(f"  총 메뉴 수:    {eval_menus}개")
    print(f"  human 쿼리:   {human_cnt}개  |  llm 쿼리: {llm_cnt}개")
    print(f"\n[인덱스 풀 (eval 제외)]")
    print(f"  top200 남은 쿼리: {len(index_pool):,}개")

    if args.dry_run:
        print("\n[dry-run] 파일 저장 안 함")
        return

    # ── 저장 ──────────────────────────────────────────────────────────────────
    with open(EVAL_OUT, "w", encoding="utf-8") as f:
        json.dump(eval_set, f, ensure_ascii=False, indent=2)
    print(f"\n평가셋 저장: {EVAL_OUT}")

    with open(INDEX_REMAINING_OUT, "w", encoding="utf-8") as f:
        json.dump(index_pool, f, ensure_ascii=False, indent=2)
    print(f"인덱스 풀 저장: {INDEX_REMAINING_OUT}")
    print("\n다음 단계: scripts/20_build_query_index.py --reset --from-pool")


if __name__ == "__main__":
    main()