from dotenv import load_dotenv

load_dotenv()  # .env 에서 OPENAI_API_KEY 로드

from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

# 검색 대상 벡터 스토어 구성 (nvme.txt + ssd.txt)
docs0 = TextLoader("../DATA/nvme.txt", encoding="utf-8").load() \
      + TextLoader("../DATA/ssd.txt",  encoding="utf-8").load()
chunks0 = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100).split_documents(docs0)
store = Chroma.from_documents(chunks0, embeddings, collection_name="retriever_demo")

# (1) 정답
retriever = store.as_retriever(search_kwargs={"k": 3})

def search_modes(store, query):
    # (2) 정답
    # 1) similarity_search — 가장 단순 / 기본
    for d in store.similarity_search(query, k=3):
        print(f"  → {d.page_content[:80]}...")

    # 2) similarity_search_with_score — 점수 포함 (Chroma 는 거리, 작을수록 유사)
    for d, score in store.similarity_search_with_score(query, k=3):
        sim_pct = round((1 - score) * 100, 1)
        print(f"  거리 {score:.3f} (유사도 ≈ {sim_pct}%)  {d.page_content[:60]}...")

    # 3) MMR — 다양성 보장
    # 같은 의미의 비슷한 청크들이 top-k 를 다 차지하지 않게,
    # 첫 결과는 가장 관련성 높은 것, 그 다음부터는 이전 결과들과 다른 것 우선.
    for d in store.max_marginal_relevance_search(query, k=3, fetch_k=10, lambda_mult=0.5):
        print(f"  → {d.page_content[:80]}...")


# (2) 정답 호출 — 검색 모드 3가지 비교
print("[검색 모드: similarity / with_score / MMR]")
search_modes(store, "NVMe 의 속도는?")


docs = [
    Document(page_content="NVMe SSD 는 PCIe 로 약 7GB/s 의 시퀀셜 속도를 낸다.",
             metadata={"category": "nvme", "year": 2022, "tier": "high"}),
    Document(page_content="SATA SSD 는 약 0.5GB/s 로 NVMe 보다 느리지만 호환성이 좋다.",
             metadata={"category": "sata", "year": 2018, "tier": "mid"}),
    Document(page_content="HDD 는 회전 디스크 기반이라 IO 가 느린 편이다.",
             metadata={"category": "hdd", "year": 2015, "tier": "low"}),
    Document(page_content="최신 PCIe 5.0 NVMe 는 약 14GB/s 까지 도달한다.",
             metadata={"category": "nvme", "year": 2024, "tier": "high"}),
]

# (3) 정답
store = Chroma.from_documents(docs, embeddings, collection_name="meta_filter_demo")

query = "저장장치 속도"

def show(title, where=None):
    print(f"\n# {title}")
    for d in store.similarity_search(query, k=4, filter=where):
        m = d.metadata
        print(f"  [{m['category']:4s} {m['year']} {m['tier']:4s}] {d.page_content[:38]}...")

# 2) 동등 — category 가 nvme 인 것만
show("category == 'nvme'", {"category": "nvme"})

# 3) 숫자 비교 — 2020년 이후 문서만
show("year >= 2020", {"year": {"$gte": 2020}})

# 5) 복합 AND — nvme 이면서 2023년 이후
show("nvme AND year >= 2023", {"$and": [{"category": "nvme"}, {"year": {"$gte": 2023}}]})
