main.py

"""
Microsoft GraphRAG 기반 지식그래프 구축 및 질의 시스템
=====================================================
국회 속기록 등 비정형 텍스트를 GraphRAG로 인덱싱하고
Neo4j로 내보내거나 Local/Global/DRIFT 검색으로 질의합니다.

사용법:
    인덱싱:       uv run main.py index
    Neo4j 적재:   uv run main.py neo4j-export
    질의:         uv run main.py query "질문 내용"
    전체 파이프라인: uv run main.py pipeline
"""

import argparse
import logging
import os
import subprocess
import sys
import time
from pathlib import Path

import pandas as pd
from dotenv import load_dotenv
from neo4j import GraphDatabase

# ──────────────────────────────────────────────
# 환경 설정
# ──────────────────────────────────────────────
load_dotenv()


# ──────────────────────────────────────────────
# litellm 패치: vLLM 커스텀 모델 호환성
# ──────────────────────────────────────────────
# litellm은 알 수 없는 모델(vLLM 커스텀 모델 포함)에 대해 두 가지를 거부함:
#   1. embedding의 dimensions 파라미터 (text-embedding-3 외 모델에서 차단)
#   2. supports_response_schema 가 False 반환 (structured output 차단)
# vLLM은 두 가지 모두 지원하므로 패치가 필요함.
_litellm_patched = False


def _patch_litellm_for_vllm():
    """litellm의 vLLM 비호환 검증들을 패치합니다.

    패치 1 — embedding dimensions:
        litellm.utils.get_optional_params_embeddings에서 OpenAI 경로가
        모델명에 'text-embedding-3' 없으면 dimensions를 무조건 거부.
        litellm.main도 모듈 레벨 import 하므로 양쪽 모두 패치 필요.

    패치 2 — supports_response_schema:
        litellm.supports_response_schema가 커스텀 모델에 대해 False 반환.
        community reports 생성 시 response_format이 필요한데 이를 차단함.
        graphrag_llm.completion.lite_llm_completion도 모듈 레벨 import 하므로
        양쪽 모두 패치 필요.
    """
    global _litellm_patched
    if _litellm_patched:
        return

    import litellm
    import litellm.main as litellm_main
    import litellm.utils as litellm_utils
    from litellm.exceptions import UnsupportedParamsError

    # ── 패치 1: embedding dimensions ──
    _original_embedding_params = litellm_utils.get_optional_params_embeddings

    def _patched_get_optional_params_embeddings(**kwargs):
        try:
            return _original_embedding_params(**kwargs)
        except UnsupportedParamsError as e:
            if "dimensions" in str(e):
                logger.info(
                    "litellm 패치: dimensions 파라미터 허용 (vLLM 커스텀 모델)"
                )
                result = {}
                for key in [
                    "dimensions",
                    "encoding_format",
                    "user",
                    "extra_body",
                    "extra_headers",
                ]:
                    val = kwargs.get(key)
                    if val is not None:
                        result[key] = val
                return result
            raise

    litellm_utils.get_optional_params_embeddings = (
        _patched_get_optional_params_embeddings
    )
    litellm_main.get_optional_params_embeddings = (
        _patched_get_optional_params_embeddings
    )

    # ── 패치 2: supports_response_schema ──
    _original_supports_response_schema = litellm.supports_response_schema

    def _patched_supports_response_schema(
        model: str, custom_llm_provider=None
    ) -> bool:
        result = _original_supports_response_schema(
            model=model, custom_llm_provider=custom_llm_provider
        )
        if not result:
            # vLLM은 guided decoding으로 structured output 지원
            logger.info(
                "litellm 패치: response_schema 지원 허용 (model=%s)", model
            )
            return True
        return result

    # litellm 모듈 자체의 함수 교체
    litellm.supports_response_schema = _patched_supports_response_schema
    litellm_utils.supports_response_schema = _patched_supports_response_schema

    # graphrag_llm이 모듈 레벨에서 import한 참조도 교체
    try:
        import graphrag_llm.completion.lite_llm_completion as _graphrag_completion

        _graphrag_completion.supports_response_schema = (
            _patched_supports_response_schema
        )
    except ImportError:
        pass

    _litellm_patched = True
    logger.info(
        "litellm vLLM 호환 패치 적용 완료 (dimensions + response_schema)"
    )

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("graphrag_app")

ROOT_DIR = Path(__file__).parent
OUTPUT_DIR = ROOT_DIR / "output"


# ══════════════════════════════════════════════
# 1. 인덱싱
# ══════════════════════════════════════════════
def run_index(verbose: bool = False) -> None:
    """GraphRAG 인덱싱을 실행합니다.

    subprocess 대신 graphrag Python API를 직접 호출하여
    litellm 패치가 적용된 상태에서 인덱싱을 수행합니다.
    """
    logger.info("=" * 60)
    logger.info("  [1/3] Microsoft GraphRAG 인덱싱 시작")
    logger.info("=" * 60)

    input_dir = ROOT_DIR / "input"
    txt_files = list(input_dir.glob("*.txt"))
    if not txt_files:
        logger.error(f"입력 파일이 없습니다: {input_dir}")
        sys.exit(1)

    total_chars = sum(f.read_text(encoding="utf-8").__len__() for f in txt_files)
    logger.info(f"입력 파일: {len(txt_files)}개, 총 {total_chars:,}자")
    for f in txt_files:
        logger.info(f"  • {f.name}")

    # litellm 패치 적용 (vLLM dimensions + response_schema 허용)
    _patch_litellm_for_vllm()

    logger.info("graphrag Python API를 직접 호출합니다 (litellm 패치 적용)")
    logger.info("-" * 60)

    t_start = time.time()
    try:
        from graphrag.cli.index import index_cli
        from graphrag.config.enums import IndexingMethod

        index_cli(
            root_dir=ROOT_DIR,
            method=IndexingMethod.Standard,
            verbose=verbose,
            cache=True,
            dry_run=False,
            skip_validation=False,
        )
        # index_cli 는 내부에서 sys.exit()를 호출함
    except SystemExit as e:
        elapsed = time.time() - t_start
        if e.code == 0 or e.code is None:
            logger.info("=" * 60)
            logger.info(f"인덱싱 완료! (소요 시간: {elapsed:.1f}초)")
            logger.info("=" * 60)
            _print_output_summary()
        else:
            logger.error(
                f"인덱싱 실패 (exit code: {e.code}, {elapsed:.1f}초)"
            )
            logger.error("상세 로그: logs/indexing-engine.log 확인")
            sys.exit(1)
    except KeyboardInterrupt:
        logger.warning("사용자에 의해 중단됨")
        sys.exit(130)
    except Exception as e:
        elapsed = time.time() - t_start
        logger.error(f"인덱싱 중 예외 발생 ({elapsed:.1f}초): {e}")
        logger.error("상세 로그: logs/indexing-engine.log 확인")
        raise


def _print_output_summary() -> None:
    """인덱싱 결과 요약을 출력합니다."""
    if not OUTPUT_DIR.exists():
        return

    logger.info("-" * 60)
    logger.info("인덱싱 결과 요약")
    logger.info("-" * 60)

    for pf in sorted(OUTPUT_DIR.rglob("*.parquet")):
        size_kb = pf.stat().st_size / 1024
        # parquet 행 수 읽기
        try:
            row_count = len(pd.read_parquet(pf))
            logger.info(f"  📄 {pf.name} — {row_count}행, {size_kb:.1f} KB")
        except Exception:
            logger.info(f"  📄 {pf.name} — {size_kb:.1f} KB")

    for gf in sorted(OUTPUT_DIR.rglob("*.graphml")):
        size_kb = gf.stat().st_size / 1024
        logger.info(f"  🔗 {gf.name} — {size_kb:.1f} KB")


# ══════════════════════════════════════════════
# 2. Neo4j 내보내기
# ══════════════════════════════════════════════
def run_neo4j_export(clear: bool = False) -> None:
    """GraphRAG 인덱싱 결과(Parquet)를 Neo4j에 적재합니다.

    생성되는 Neo4j 그래프 구조:
        노드: (:Entity), (:Community), (:CommunityReport), (:TextUnit), (:Document)
        관계: (:Entity)-[:RELATES_TO]->(:Entity)
              (:Entity)-[:BELONGS_TO_COMMUNITY]->(:Community)
              (:Entity)-[:FOUND_IN]->(:TextUnit)
              (:TextUnit)-[:PART_OF]->(:Document)
              (:Community)-[:HAS_REPORT]->(:CommunityReport)
              (:Community)-[:PARENT_OF]->(:Community)
    """
    logger.info("=" * 60)
    logger.info("  [2/3] Neo4j 내보내기")
    logger.info("=" * 60)

    # 필수 parquet 파일 확인
    required_files = ["entities.parquet", "relationships.parquet"]
    for fname in required_files:
        if not (OUTPUT_DIR / fname).exists():
            logger.error(f"필수 파일 없음: {OUTPUT_DIR / fname}")
            logger.error("먼저 'uv run main.py index'를 실행하세요.")
            sys.exit(1)

    # Neo4j 연결
    uri = os.environ.get("NEO4J_URI", "neo4j://localhost:7687")
    user = os.environ.get("NEO4J_USERNAME", "neo4j")
    password = os.environ.get("NEO4J_PASSWORD", "")

    driver = GraphDatabase.driver(uri, auth=(user, password))
    try:
        driver.verify_connectivity()
        logger.info(f"Neo4j 연결 성공: {uri}")
    except Exception as e:
        logger.error(f"Neo4j 연결 실패: {e}")
        sys.exit(1)

    t_start = time.time()

    try:
        with driver.session() as session:
            # 기존 데이터 정리
            if clear:
                logger.info("기존 데이터 삭제 중...")
                session.run("MATCH (n) DETACH DELETE n")
                logger.info("  ✅ 기존 데이터 삭제 완료")

            # 인덱스 생성
            _create_indexes(session)

            # 데이터 적재
            _load_entities(session)
            _load_relationships(session)
            _load_communities(session)
            _load_community_reports(session)
            _load_text_units(session)
            _load_documents(session)

            # 관계 연결
            _link_entities_to_communities(session)
            _link_entities_to_text_units(session)
            _link_text_units_to_documents(session)
            _link_community_hierarchy(session)
            _link_communities_to_reports(session)

        elapsed = time.time() - t_start
        logger.info("=" * 60)
        logger.info(f"Neo4j 내보내기 완료! (소요 시간: {elapsed:.1f}초)")
        logger.info("=" * 60)

        _print_neo4j_summary(driver)

    finally:
        driver.close()


def _create_indexes(session) -> None:
    """Neo4j 인덱스를 생성합니다."""
    logger.info("인덱스 생성 중...")
    indexes = [
        "CREATE INDEX IF NOT EXISTS FOR (e:Entity) ON (e.id)",
        "CREATE INDEX IF NOT EXISTS FOR (e:Entity) ON (e.title)",
        "CREATE INDEX IF NOT EXISTS FOR (c:Community) ON (c.id)",
        "CREATE INDEX IF NOT EXISTS FOR (c:Community) ON (c.community)",
        "CREATE INDEX IF NOT EXISTS FOR (r:CommunityReport) ON (r.community)",
        "CREATE INDEX IF NOT EXISTS FOR (t:TextUnit) ON (t.id)",
        "CREATE INDEX IF NOT EXISTS FOR (d:Document) ON (d.id)",
    ]
    for idx in indexes:
        session.run(idx)
    logger.info(f"  ✅ 인덱스 {len(indexes)}개 생성/확인")


def _load_entities(session) -> None:
    """entities.parquet → (:Entity) 노드"""
    path = OUTPUT_DIR / "entities.parquet"
    if not path.exists():
        return

    df = pd.read_parquet(path)
    logger.info(f"엔티티 로드 중... ({len(df)}개)")

    records = []
    for _, row in df.iterrows():
        records.append({
            "id": str(row.get("id", "")),
            "title": str(row.get("title", "")),
            "type": str(row.get("type", "")),
            "description": str(row.get("description", "")),
            "human_readable_id": int(row["human_readable_id"]) if pd.notna(row.get("human_readable_id")) else 0,
            "frequency": int(row["frequency"]) if pd.notna(row.get("frequency")) else 0,
            "degree": int(row["degree"]) if pd.notna(row.get("degree")) else 0,
        })

    session.run(
        """
        UNWIND $records AS rec
        MERGE (e:Entity {id: rec.id})
        SET e.title = rec.title,
            e.type = rec.type,
            e.description = rec.description,
            e.human_readable_id = rec.human_readable_id,
            e.frequency = rec.frequency,
            e.degree = rec.degree
        """,
        records=records,
    )

    # 엔티티 타입별 라벨 추가 (Person, Organization 등)
    session.run(
        """
        MATCH (e:Entity)
        WHERE e.type IS NOT NULL AND e.type <> ''
        WITH e, apoc.text.capitalize(e.type) AS label
        CALL apoc.create.addLabels(e, [label]) YIELD node
        RETURN count(node)
        """
    )

    logger.info(f"  ✅ Entity 노드 {len(records)}개 생성")


def _load_relationships(session) -> None:
    """relationships.parquet → (:Entity)-[:RELATES_TO]->(:Entity)"""
    path = OUTPUT_DIR / "relationships.parquet"
    if not path.exists():
        return

    df = pd.read_parquet(path)
    logger.info(f"관계 로드 중... ({len(df)}개)")

    records = []
    for _, row in df.iterrows():
        records.append({
            "id": str(row.get("id", "")),
            "source": str(row.get("source", "")),
            "target": str(row.get("target", "")),
            "description": str(row.get("description", "")),
            "weight": float(row["weight"]) if pd.notna(row.get("weight")) else 1.0,
            "combined_degree": int(row["combined_degree"]) if pd.notna(row.get("combined_degree")) else 0,
            "human_readable_id": int(row["human_readable_id"]) if pd.notna(row.get("human_readable_id")) else 0,
        })

    session.run(
        """
        UNWIND $records AS rec
        MATCH (src:Entity {title: rec.source})
        MATCH (tgt:Entity {title: rec.target})
        MERGE (src)-[r:RELATES_TO {id: rec.id}]->(tgt)
        SET r.description = rec.description,
            r.weight = rec.weight,
            r.combined_degree = rec.combined_degree,
            r.human_readable_id = rec.human_readable_id
        """,
        records=records,
    )
    logger.info(f"  ✅ RELATES_TO 관계 {len(records)}개 생성")


def _load_communities(session) -> None:
    """communities.parquet → (:Community) 노드"""
    path = OUTPUT_DIR / "communities.parquet"
    if not path.exists():
        logger.info("  ⏭️  communities.parquet 없음, 건너뜀")
        return

    df = pd.read_parquet(path)
    logger.info(f"커뮤니티 로드 중... ({len(df)}개)")

    records = []
    for _, row in df.iterrows():
        records.append({
            "id": str(row.get("id", "")),
            "community": int(row["community"]) if pd.notna(row.get("community")) else 0,
            "level": int(row["level"]) if pd.notna(row.get("level")) else 0,
            "title": str(row.get("title", "")),
            "parent": int(row["parent"]) if pd.notna(row.get("parent")) else -1,
            "size": int(row["size"]) if pd.notna(row.get("size")) else 0,
        })

    session.run(
        """
        UNWIND $records AS rec
        MERGE (c:Community {id: rec.id})
        SET c.community = rec.community,
            c.level = rec.level,
            c.title = rec.title,
            c.parent = rec.parent,
            c.size = rec.size
        """,
        records=records,
    )
    logger.info(f"  ✅ Community 노드 {len(records)}개 생성")


def _load_community_reports(session) -> None:
    """community_reports.parquet → (:CommunityReport) 노드"""
    path = OUTPUT_DIR / "community_reports.parquet"
    if not path.exists():
        logger.info("  ⏭️  community_reports.parquet 없음, 건너뜀")
        return

    df = pd.read_parquet(path)
    logger.info(f"커뮤니티 리포트 로드 중... ({len(df)}개)")

    records = []
    for _, row in df.iterrows():
        records.append({
            "id": str(row.get("id", "")),
            "community": int(row["community"]) if pd.notna(row.get("community")) else 0,
            "level": int(row["level"]) if pd.notna(row.get("level")) else 0,
            "title": str(row.get("title", "")),
            "summary": str(row.get("summary", "")),
            "full_content": str(row.get("full_content", ""))[:5000],  # 너무 긴 텍스트 제한
            "rank": float(row["rank"]) if pd.notna(row.get("rank")) else 0.0,
            "rating_explanation": str(row.get("rating_explanation", "")),
        })

    session.run(
        """
        UNWIND $records AS rec
        MERGE (cr:CommunityReport {id: rec.id})
        SET cr.community = rec.community,
            cr.level = rec.level,
            cr.title = rec.title,
            cr.summary = rec.summary,
            cr.full_content = rec.full_content,
            cr.rank = rec.rank,
            cr.rating_explanation = rec.rating_explanation
        """,
        records=records,
    )
    logger.info(f"  ✅ CommunityReport 노드 {len(records)}개 생성")


def _load_text_units(session) -> None:
    """text_units.parquet → (:TextUnit) 노드"""
    path = OUTPUT_DIR / "text_units.parquet"
    if not path.exists():
        logger.info("  ⏭️  text_units.parquet 없음, 건너뜀")
        return

    df = pd.read_parquet(path)
    logger.info(f"텍스트 유닛 로드 중... ({len(df)}개)")

    records = []
    for _, row in df.iterrows():
        records.append({
            "id": str(row.get("id", "")),
            "text": str(row.get("text", ""))[:2000],  # 텍스트 길이 제한
            "n_tokens": int(row["n_tokens"]) if pd.notna(row.get("n_tokens")) else 0,
            "document_id": str(row.get("document_id", "")),
        })

    session.run(
        """
        UNWIND $records AS rec
        MERGE (t:TextUnit {id: rec.id})
        SET t.text = rec.text,
            t.n_tokens = rec.n_tokens,
            t.document_id = rec.document_id
        """,
        records=records,
    )
    logger.info(f"  ✅ TextUnit 노드 {len(records)}개 생성")


def _load_documents(session) -> None:
    """documents.parquet → (:Document) 노드"""
    path = OUTPUT_DIR / "documents.parquet"
    if not path.exists():
        logger.info("  ⏭️  documents.parquet 없음, 건너뜀")
        return

    df = pd.read_parquet(path)
    logger.info(f"문서 로드 중... ({len(df)}개)")

    records = []
    for _, row in df.iterrows():
        records.append({
            "id": str(row.get("id", "")),
            "title": str(row.get("title", "")),
        })

    session.run(
        """
        UNWIND $records AS rec
        MERGE (d:Document {id: rec.id})
        SET d.title = rec.title
        """,
        records=records,
    )
    logger.info(f"  ✅ Document 노드 {len(records)}개 생성")


# ── 관계 연결 ──

def _link_entities_to_communities(session) -> None:
    """엔티티 → 커뮤니티 관계 연결"""
    path_c = OUTPUT_DIR / "communities.parquet"
    if not path_c.exists():
        return

    df = pd.read_parquet(path_c)
    logger.info("엔티티-커뮤니티 관계 연결 중...")

    link_count = 0
    for _, row in df.iterrows():
        entity_ids = row.get("entity_ids", [])
        if entity_ids is None or (isinstance(entity_ids, float) and pd.isna(entity_ids)):
            continue
        community_id = str(row.get("id", ""))

        records = [{"entity_id": str(eid), "community_id": community_id} for eid in entity_ids]
        if records:
            session.run(
                """
                UNWIND $records AS rec
                MATCH (e:Entity {id: rec.entity_id})
                MATCH (c:Community {id: rec.community_id})
                MERGE (e)-[:BELONGS_TO_COMMUNITY]->(c)
                """,
                records=records,
            )
            link_count += len(records)

    logger.info(f"  ✅ BELONGS_TO_COMMUNITY 관계 {link_count}개 생성")


def _link_entities_to_text_units(session) -> None:
    """엔티티 → 텍스트유닛 관계 연결"""
    path = OUTPUT_DIR / "entities.parquet"
    if not path.exists():
        return

    df = pd.read_parquet(path)
    logger.info("엔티티-텍스트유닛 관계 연결 중...")

    link_count = 0
    for _, row in df.iterrows():
        text_unit_ids = row.get("text_unit_ids", [])
        if text_unit_ids is None or (isinstance(text_unit_ids, float) and pd.isna(text_unit_ids)):
            continue
        entity_id = str(row.get("id", ""))

        records = [{"entity_id": entity_id, "tu_id": str(tid)} for tid in text_unit_ids]
        if records:
            session.run(
                """
                UNWIND $records AS rec
                MATCH (e:Entity {id: rec.entity_id})
                MATCH (t:TextUnit {id: rec.tu_id})
                MERGE (e)-[:FOUND_IN]->(t)
                """,
                records=records,
            )
            link_count += len(records)

    logger.info(f"  ✅ FOUND_IN 관계 {link_count}개 생성")


def _link_text_units_to_documents(session) -> None:
    """텍스트유닛 → 문서 관계 연결"""
    session.run(
        """
        MATCH (t:TextUnit)
        WHERE t.document_id IS NOT NULL AND t.document_id <> ''
        MATCH (d:Document {id: t.document_id})
        MERGE (t)-[:PART_OF]->(d)
        """
    )
    logger.info("  ✅ PART_OF (TextUnit→Document) 관계 생성")


def _link_community_hierarchy(session) -> None:
    """커뮤니티 계층 관계 연결"""
    path = OUTPUT_DIR / "communities.parquet"
    if not path.exists():
        return

    df = pd.read_parquet(path)
    logger.info("커뮤니티 계층 관계 연결 중...")

    records = []
    for _, row in df.iterrows():
        children = row.get("children", [])
        if children is None or (isinstance(children, float) and pd.isna(children)):
            continue
        parent_id = str(row.get("id", ""))
        for child_id in children:
            records.append({"parent_id": parent_id, "child_id": str(child_id)})

    if records:
        session.run(
            """
            UNWIND $records AS rec
            MATCH (parent:Community {id: rec.parent_id})
            MATCH (child:Community {community: toInteger(rec.child_id)})
            MERGE (parent)-[:PARENT_OF]->(child)
            """,
            records=records,
        )
    logger.info(f"  ✅ PARENT_OF 관계 {len(records)}개 생성")


def _link_communities_to_reports(session) -> None:
    """커뮤니티 → 리포트 관계 연결"""
    session.run(
        """
        MATCH (c:Community)
        MATCH (cr:CommunityReport {community: c.community})
        MERGE (c)-[:HAS_REPORT]->(cr)
        """
    )
    logger.info("  ✅ HAS_REPORT (Community→CommunityReport) 관계 생성")


def _print_neo4j_summary(driver) -> None:
    """Neo4j 저장 결과 요약"""
    logger.info("-" * 60)
    logger.info("Neo4j 저장 결과 요약")
    logger.info("-" * 60)

    with driver.session() as session:
        # 노드 현황
        label_counts = session.run(
            "CALL db.labels() YIELD label "
            "RETURN label, count { MATCH (n) WHERE label IN labels(n) RETURN count(n) } AS cnt "
            "ORDER BY cnt DESC"
        ).data()

        if label_counts:
            logger.info("[노드 현황]")
            for row in label_counts:
                logger.info(f"  • :{row['label']}  →  {row['cnt']}개")

        # 관계 현황
        rel_counts = session.run(
            "CALL db.relationshipTypes() YIELD relationshipType AS type "
            "RETURN type, count { MATCH ()-[r]->() WHERE type(r) = type RETURN count(r) } AS cnt "
            "ORDER BY cnt DESC"
        ).data()

        if rel_counts:
            logger.info("[관계 현황]")
            for row in rel_counts:
                logger.info(f"  • {row['type']}  →  {row['cnt']}개")

        # 엔티티 샘플
        samples = session.run(
            "MATCH (e:Entity) "
            "RETURN e.title AS title, e.type AS type, e.description AS desc "
            "ORDER BY e.degree DESC "
            "LIMIT 10"
        ).data()

        if samples:
            logger.info("[주요 엔티티 (연결도 상위 10)]")
            for row in samples:
                desc = (row["desc"] or "")[:80]
                logger.info(f"  • [{row['type']}] {row['title']} — {desc}...")


# ══════════════════════════════════════════════
# 3. 질의
# ══════════════════════════════════════════════
def run_query(question: str, method: str = "local") -> None:
    """GraphRAG 질의를 실행합니다."""
    valid_methods = ["local", "global", "drift", "basic"]
    if method not in valid_methods:
        logger.error(f"잘못된 검색 방법: {method} (가능: {valid_methods})")
        sys.exit(1)

    logger.info("=" * 60)
    logger.info(f"  [3/3] GraphRAG 질의 ({method} search)")
    logger.info(f"  Q: {question}")
    logger.info("=" * 60)

    if not OUTPUT_DIR.exists() or not list(OUTPUT_DIR.rglob("*.parquet")):
        logger.error("인덱싱 결과가 없습니다. 먼저 'uv run main.py index'를 실행하세요.")
        sys.exit(1)

    cmd = [
        "uv", "run", "graphrag", "query",
        "--root", str(ROOT_DIR),
        "--method", method,
        "--query", question,
    ]

    logger.info(f"실행 명령: {' '.join(cmd)}")
    logger.info("-" * 60)

    t_start = time.time()
    try:
        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            cwd=str(ROOT_DIR),
        )
        elapsed = time.time() - t_start

        if result.returncode == 0:
            print("\n" + result.stdout)
            logger.info(f"질의 완료 ({elapsed:.1f}초)")
        else:
            logger.error(f"질의 실패 (exit code: {result.returncode})")
            if result.stderr:
                print(result.stderr)
            sys.exit(1)

    except KeyboardInterrupt:
        logger.warning("사용자에 의해 중단됨")
        sys.exit(130)


# ══════════════════════════════════════════════
# 전체 파이프라인 (index → neo4j-export)
# ══════════════════════════════════════════════
def run_pipeline(verbose: bool = False, clear: bool = False) -> None:
    """인덱싱 → Neo4j 적재를 한 번에 실행합니다."""
    run_index(verbose=verbose)
    run_neo4j_export(clear=clear)
    logger.info("=" * 60)
    logger.info("  전체 파이프라인 완료!")
    logger.info("  이제 'uv run main.py query \"질문\"'으로 질의할 수 있습니다.")
    logger.info("=" * 60)


# ══════════════════════════════════════════════
# CLI
# ══════════════════════════════════════════════
def main() -> None:
    parser = argparse.ArgumentParser(
        description="Microsoft GraphRAG 기반 지식그래프 구축 및 질의 시스템",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
사용 예시:
  uv run main.py index                                        # 인덱싱만 실행
  uv run main.py neo4j-export                                 # Neo4j로 적재
  uv run main.py neo4j-export --clear                         # 기존 데이터 삭제 후 적재
  uv run main.py pipeline --clear                             # 인덱싱 + Neo4j 적재
  uv run main.py query "검찰개혁 청문회에서 핵심 쟁점은?"              # Local 검색
  uv run main.py query --method global "전체 회의 분위기는?"         # Global 검색
        """,
    )

    subparsers = parser.add_subparsers(dest="command", help="실행할 명령")

    # index
    index_parser = subparsers.add_parser("index", help="GraphRAG 인덱싱 실행")
    index_parser.add_argument("--verbose", "-v", action="store_true", help="상세 로그 출력")

    # neo4j-export
    export_parser = subparsers.add_parser("neo4j-export", help="인덱싱 결과를 Neo4j에 적재")
    export_parser.add_argument("--clear", action="store_true", help="기존 Neo4j 데이터 삭제 후 적재")

    # pipeline (index + neo4j-export)
    pipe_parser = subparsers.add_parser("pipeline", help="인덱싱 + Neo4j 적재를 한 번에 실행")
    pipe_parser.add_argument("--verbose", "-v", action="store_true", help="상세 로그 출력")
    pipe_parser.add_argument("--clear", action="store_true", help="기존 Neo4j 데이터 삭제 후 적재")

    # query
    query_parser = subparsers.add_parser("query", help="GraphRAG 질의 실행")
    query_parser.add_argument("question", type=str, help="질문 내용")
    query_parser.add_argument(
        "--method", "-m",
        type=str,
        default="local",
        choices=["local", "global", "drift", "basic"],
        help="검색 방법 (기본: local)",
    )

    args = parser.parse_args()

    if args.command == "index":
        run_index(verbose=args.verbose)
    elif args.command == "neo4j-export":
        run_neo4j_export(clear=args.clear)
    elif args.command == "pipeline":
        run_pipeline(verbose=args.verbose, clear=args.clear)
    elif args.command == "query":
        run_query(args.question, method=args.method)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()