scripts/20_build_query_index.py

"""
쿼리 인덱스 구축 스크립트
=========================
엑셀(MenuSearch_review_all_20260521.xlsx)에서 현업 검증 쿼리를 추출해
ChromaDB "herogun_queries" 컬렉션에 저장.

현업 작업 내용 반영:
  - sample_queries(LLM 생성): 취소선 없는 것만 사용 (취소선 = 나쁜 쿼리 표시)
  - Sample_queries(자체 생성): 현업이 직접 추가한 쿼리 (전부 사용)
  - 미검토 메뉴(취소선/자체생성 없음): LLM 쿼리 전체 사용 (검토 안된 것)

소스 구분:
  "human"         - 현업 직접 작성 (자체 생성)
  "llm_reviewed"  - 현업이 취소선 검토 후 유지한 LLM 쿼리
  "llm_unreviewed"- 미검토 메뉴의 LLM 쿼리 (노이즈 가능성 있음)

실행:
    .venv/Scripts/python.exe scripts/20_build_query_index.py
    .venv/Scripts/python.exe scripts/20_build_query_index.py --reset
    .venv/Scripts/python.exe scripts/20_build_query_index.py --reviewed-only  # 검토된 메뉴만
    .venv/Scripts/python.exe scripts/20_build_query_index.py --stats
"""
import sys
import re
import argparse
from pathlib import Path

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

from openpyxl import load_workbook
from openpyxl.cell.rich_text import CellRichText, TextBlock
import pandas as pd
from tqdm import tqdm

from core.embedder import MenuEmbedder
from core.query_vectorstore import QueryVectorStore

EXCEL_PATH   = Path(__file__).parent.parent / "data" / "MenuSearch_review_all_20260525_edit.xlsx"
POOL_PATH    = Path(__file__).parent.parent / "data" / "query_index_pool.json"  # eval 제외 후 남은 풀
BATCH_SIZE   = 64
TOP_N_CUTOFF = 200  # 현업 검토 범위: top200까지만 인덱싱 (201+ 미검토 → 노이즈 제외)

LLM_COL   = 8  # sample_queries(LLM 생성)
HUMAN_COL = 9  # Sample_queries(자체 생성)


def parse_text_to_queries(text: str) -> list[str]:
    """번호 리스트 또는 줄바꿈으로 구분된 텍스트를 쿼리 목록으로 파싱"""
    if not text or text.strip() in ("", "nan"):
        return []
    text = text.strip()
    # 번호 리스트 "1. " "2) " 패턴
    parts = re.split(r"\n?\s*\d+[.)]\s+", text)
    if len(parts) > 1:
        queries = [p.strip() for p in parts if p.strip()]
    else:
        queries = [p.strip() for p in re.split(r"\s*/\s*|\n", text) if p.strip()]
    return [q for q in queries if len(q) >= 4 and not re.match(r"^(예:|N/A|없음|해당없음)", q)]


def extract_rich_text_parts(val) -> tuple[list[str], list[str]]:
    """
    rich_text 셀에서 (유지 텍스트 목록, 취소선 텍스트 목록) 반환.
    plain text면 (전체 쿼리 목록, []) 반환.
    """
    if val is None:
        return [], []

    if isinstance(val, CellRichText):
        kept_raw = []
        removed_raw = []
        for block in val:
            if isinstance(block, TextBlock):
                text = str(block.text)
                if block.font and block.font.strike:
                    removed_raw.append(text)
                else:
                    kept_raw.append(text)
            elif isinstance(block, str):
                kept_raw.append(block)

        kept   = parse_text_to_queries("".join(kept_raw))
        removed = parse_text_to_queries("".join(removed_raw))
        return kept, removed
    else:
        # 일반 텍스트
        return parse_text_to_queries(str(val)), []


def load_queries_from_excel(excel_path: Path, reviewed_only: bool = False) -> list[dict]:
    """
    엑셀에서 쿼리 추출.

    Returns:
        list of {
            "query":     str,
            "menu_id":   str,
            "menu_name": str,
            "menu_path": str,
            "source":    "human" | "llm_reviewed" | "llm_unreviewed",
        }
    """
    print(f"엑셀 로드: {excel_path}")
    wb = load_workbook(excel_path, rich_text=True)
    ws = wb["재생성메뉴검토"]

    # pandas로 menu_id/name/path 읽기
    df = pd.read_excel(excel_path, sheet_name="재생성메뉴검토", engine="openpyxl")

    rows = []
    counts = {"human": 0, "llm_top200": 0, "skipped_rank": 0, "removed": 0}

    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

        # top200 이하만 사용 (현업 검토 범위)
        try:
            rank_int = int(rank)
        except (TypeError, ValueError):
            rank_int = 9999

        if rank_int > TOP_N_CUTOFF:
            counts["skipped_rank"] += 1
            continue

        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 쿼리 (취소선 제외) ───────────────────────────────────────
        val_llm = ws.cell(row_idx, LLM_COL).value
        kept_llm, removed_llm = extract_rich_text_parts(val_llm)
        counts["removed"] += len(removed_llm)

        # ── 현업 자체 추가 쿼리 ─────────────────────────────────────────
        val_human = ws.cell(row_idx, HUMAN_COL).value
        human_queries, _ = extract_rich_text_parts(val_human)

        if reviewed_only:
            is_reviewed = (len(removed_llm) > 0) or (len(human_queries) > 0)
            if not is_reviewed:
                continue

        for q in human_queries:
            rows.append({
                "query":     q,
                "menu_id":   menu_id,
                "menu_name": menu_name,
                "menu_path": menu_path,
                "source":    "human",
            })
            counts["human"] += 1

        for q in kept_llm:
            rows.append({
                "query":     q,
                "menu_id":   menu_id,
                "menu_name": menu_name,
                "menu_path": menu_path,
                "source":    "llm_top200",
            })
            counts["llm_top200"] += 1

    print(f"\n  추출 결과 (top{TOP_N_CUTOFF} 기준):")
    print(f"    현업 직접 작성 (human):    {counts['human']:>5}개")
    print(f"    LLM top200 (취소선 제외):  {counts['llm_top200']:>5}개")
    print(f"    취소선으로 제거됨:          {counts['removed']:>5}개")
    print(f"    201위 이하 스킵:            {counts['skipped_rank']:>5}개 메뉴")
    print(f"    ──────────────────────────────────")
    print(f"    인덱싱 총합:               {len(rows):>5}개")
    return rows


def build_query_index(rows: list[dict], reset: bool = False):
    store = QueryVectorStore()
    if reset:
        store.reset()
        print(f"\n기존 컬렉션 '{store.collection_name}' 초기화")
    else:
        store.ensure_collection()

    embedder = MenuEmbedder()

    print(f"\n임베딩 및 저장 ({len(rows):,}개)...")
    for start in tqdm(range(0, len(rows), BATCH_SIZE), desc="저장"):
        batch = rows[start:start + BATCH_SIZE]
        embeddings = embedder.embed_batch([r["query"] for r in batch])
        store.upsert_queries(batch, embeddings)

    final = store.count()
    print(f"\n완료 → 컬렉션 총 {final:,}개 (Qdrant, 영속화 자동)")
    return store


def print_stats():
    store = QueryVectorStore()
    if not store.available:
        print(f"컬렉션 '{store.collection_name}' 없음. 먼저 인덱스를 구축하세요.")
        return
    total = store.count()
    print(f"\n[쿼리 인덱스 통계]  컬렉션: {store.collection_name}")
    print(f"  총 쿼리 수: {total:,}")
    # source 분포 (Qdrant scroll)
    sources: dict[str, int] = {}
    offset = None
    while True:
        points, offset = store.client.scroll(
            collection_name=store.collection_name,
            limit=1000, offset=offset, with_payload=True, with_vectors=False,
        )
        for p in points:
            src = (p.payload or {}).get("source", "unknown")
            sources[src] = sources.get(src, 0) + 1
        if offset is None:
            break
    for src, cnt in sorted(sources.items()):
        print(f"  - {src}: {cnt:,}개")


def main():
    parser = argparse.ArgumentParser(description="쿼리 인덱스 구축")
    parser.add_argument("--reset",         action="store_true", help="기존 컬렉션 삭제 후 재구축")
    parser.add_argument("--reviewed-only", action="store_true", help="현업 검토 메뉴만 (미검토 LLM 제외)")
    parser.add_argument("--from-pool",     action="store_true", help="eval 제외 풀(query_index_pool.json) 사용")
    parser.add_argument("--stats",         action="store_true", help="통계만 출력")
    args = parser.parse_args()

    if args.stats:
        print_stats()
        return

    if args.from_pool:
        # eval300 제외한 나머지 풀 사용 (데이터 누수 방지)
        import json
        if not POOL_PATH.exists():
            print(f"ERROR: {POOL_PATH} 없음. 먼저 scripts/21_build_eval300.py 실행하세요.")
            return
        with open(POOL_PATH, "r", encoding="utf-8") as f:
            pool = json.load(f)
        rows = [{
            "query":     x["query"],
            "menu_id":   x["menu_id"],
            "menu_name": x["menu_name"],
            "menu_path": x["menu_path"],
            "source":    "human" if x["source"] == "human" else "llm_top200",
        } for x in pool]
        print(f"eval 제외 풀: {len(rows):,}개")
    else:
        rows = load_queries_from_excel(EXCEL_PATH, reviewed_only=args.reviewed_only)

    if not rows:
        print("ERROR: 추출된 쿼리가 없습니다.")
        return

    build_query_index(rows, reset=args.reset)
    print_stats()


if __name__ == "__main__":
    main()