# -*- coding: utf-8 -*-
"""
build_rag.py

기능:
- chunk_and_summary.py 에서 생성된 output/rag/*_chunks.json 파일들을 읽어서
  text + summary 를 임베딩(text-embedding-3-small)한다.
- FAISS IndexFlatIP 인덱스를 구축하여
  output/rag/faiss.index, meta.json, vectors.npy 를 생성한다.
"""

import os
import sys
import json
from pathlib import Path

import numpy as np
import faiss
from openai import OpenAI
from api_config import API_KEYS

# ===== 경로 설정 =====
DATA_ROOT   = Path(r"D:\for python\survey_test\process")
OUTPUT_ROOT = Path(r"D:\for python\survey_test\output")
RAG_DIR     = OUTPUT_ROOT / "rag"
LOG_DIR     = OUTPUT_ROOT / "logs"

for d in [RAG_DIR, LOG_DIR]:
    d.mkdir(parents=True, exist_ok=True)

# ===== OpenAI 설정 (구조 유지) =====
OPENAI_API_KEY = API_KEYS.get('GPT_API_KEY', '')
GPT_MODEL      = "gpt-5-2025-08-07"
EMBED_MODEL    = "text-embedding-3-small"

client = OpenAI(api_key=OPENAI_API_KEY)


def log(msg: str):
    print(msg, flush=True)
    with (LOG_DIR / "build_rag_log.txt").open("a", encoding="utf-8") as f:
        f.write(msg + "\n")


def embed_texts(texts):
    if not texts:
        return np.zeros((0, 1536), dtype="float32")
    embs = []
    B = 96
    for i in range(0, len(texts), B):
        batch = texts[i:i+B]
        resp = client.embeddings.create(model=EMBED_MODEL, input=batch)
        for d in resp.data:
            embs.append(np.array(d.embedding, dtype="float32"))
    return np.vstack(embs)


def _build_embed_input(u: dict) -> str:
    """
    text + summary 를 합쳐 임베딩 입력을 만든다.
    - text, summary 중 없는 것은 생략
    - 공백 정리
    - 최대 길이 제한
    """
    sum_ = (u.get("summary") or "").strip()
    txt  = (u.get("text") or "").strip()

    if txt and sum_:
        merged = txt + "\n\n요약: " + sum_[:1000]
    else:
        merged = txt or sum_

    merged = " ".join(merged.split())
    if not merged:
        return ""
    if len(merged) > 4000:
        merged = merged[:4000]
    return merged


def build_faiss_index():
    docs = []
    metas = []

    rag_files = list(RAG_DIR.glob("*_chunks.json"))
    if not rag_files:
        log("RAG 파일(*_chunks.json)이 없습니다. 먼저 chunk_and_summary.py를 실행해야 합니다.")
        sys.exit(1)

    for f in rag_files:
        try:
            units = json.loads(f.read_text(encoding="utf-8", errors="ignore"))
        except Exception as e:
            log(f"[WARN] RAG 파일 읽기 실패: {f.name} | {e}")
            continue

        for u in units:
            embed_input = _build_embed_input(u)
            if not embed_input:
                continue
            if len(embed_input) < 40:
                continue
            docs.append(embed_input)
            metas.append({
                "source": u.get("source", ""),
                "chunk": int(u.get("chunk", 0)),
                "folder_context": u.get("folder_context", "")
            })

    if not docs:
        log("임베딩할 텍스트가 없습니다.")
        sys.exit(1)

    log(f"임베딩 대상 텍스트 수: {len(docs)}")

    E = embed_texts(docs)
    if E.shape[0] != len(docs):
        log(f"[WARN] 임베딩 수 불일치: E={E.shape[0]}, docs={len(docs)}")

    faiss.normalize_L2(E)
    index = faiss.IndexFlatIP(E.shape[1])
    index.add(E)

    np.save(str(RAG_DIR / "vectors.npy"), E)
    (RAG_DIR / "meta.json").write_text(
        json.dumps(metas, ensure_ascii=False, indent=2),
        encoding="utf-8"
    )
    faiss.write_index(index, str(RAG_DIR / "faiss.index"))

    log(f"FAISS 인덱스 구축 완료: 벡터 수={len(metas)}")


def main():
    log("=== FAISS RAG 인덱스 구축 시작 ===")
    build_faiss_index()
    log("=== FAISS RAG 인덱스 구축 종료 ===")


if __name__ == "__main__":
    main()