from dotenv import load_dotenv

load_dotenv()  # .env 에서 OPENAI_API_KEY 로드

import os
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
FAISS_DIR = "./faiss_index"

def build_index():
    # (1) 정답
    docs = TextLoader("../DATA/nvme.txt", encoding="utf-8").load() \
         + TextLoader("../DATA/ssd.txt",  encoding="utf-8").load()
    chunks = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100).split_documents(docs)
    for c in chunks:
        c.metadata["source"] = os.path.basename(c.metadata.get("source", "?"))

    store = FAISS.from_documents(chunks, embeddings)

    # ★ 차이 1: 영속화는 명시적 호출. Chroma 는 persist_directory 만 주면 자동.
    store.save_local(FAISS_DIR)
    return store

def load_index():
    # (2) 정답
    # ★ 차이 2: pickle 역직렬화이므로 allow_dangerous_deserialization 명시 필요.
    store = FAISS.load_local(
        FAISS_DIR,
        embeddings,
        allow_dangerous_deserialization=True,
    )
    return store

def search_demo(store):
    # (3) 정답
    # (A) FAISS 가 지원하는 가벼운 filter
    docs_a = store.similarity_search("PCIe 인터페이스", k=3, filter={"source": "nvme.txt"})

    # (B) 복잡한 조건은 직접 후처리
    all_docs = store.similarity_search("PCIe", k=10)
    docs_b = [d for d in all_docs if d.metadata.get("source") in {"nvme.txt", "ssd.txt"}][:3]
    return docs_a, docs_b


# 실제 실행: 인덱스 생성 → 검색
store = build_index()
docs_a, docs_b = search_demo(store)
print('[A] filter={"source": "nvme.txt"} 상위 결과')
for d in docs_a:
    print(f"  {d.page_content[:45]}... ({d.metadata['source']})")
print('\n[B] 후처리 필터 (nvme.txt | ssd.txt) 상위')
for d in docs_b:
    print(f"  {d.page_content[:45]}... ({d.metadata['source']})")
