from dotenv import load_dotenv

load_dotenv()  # .env 에서 OPENAI_API_KEY 로드

import os
from langchain_community.document_loaders import TextLoader, PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma

PERSIST_DIR     = "./chroma_db"
COLLECTION_NAME = "storage_demo"
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

def build_store():
    # (1) 정답
    """문서 로드 → 청킹 → 임베딩 → Chroma 에 저장"""
    docs   = TextLoader("../DATA/nvme.txt", encoding="utf-8").load()
    chunks = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100).split_documents(docs)

    store = Chroma.from_documents(
        chunks, embeddings,
        collection_name=COLLECTION_NAME,
        persist_directory=PERSIST_DIR,
    )
    print(f"새 DB 생성 — {len(chunks)} 청크 임베딩 저장")
    return store

def load_store():
    # (2) 정답
    """디스크에서 기존 DB 로드 (임베딩 재계산 없음)"""
    store = Chroma(
        collection_name=COLLECTION_NAME,
        embedding_function=embeddings,
        persist_directory=PERSIST_DIR,
    )
    print(f"기존 DB 로드 — {store._collection.count()} 청크 있음")
    return store

# (2) 정답 ("있으면 로드, 없으면 생성" 분기)
# 디스크에 컬렉션이 이미 있으면 로드, 없으면 새로 만들기
if os.path.exists(PERSIST_DIR) and os.listdir(PERSIST_DIR):
    store = load_store()
else:
    store = build_store()

def inspect_store(store):
    # (3) 정답
    # 1) 전체 문서 개수
    count = store._collection.count()
    print(f"컬렉션 '{COLLECTION_NAME}' 의 문서 수: {count}\n")

    # 2) 일부 문서 샘플 — get(limit=N) 으로 그냥 가져오기 (검색 아님)
    results = store._collection.get(include=["documents", "metadatas"], limit=3)
    for i, (doc_id, content, meta) in enumerate(
        zip(results["ids"], results["documents"], results["metadatas"]), start=1):
        print(f"[{i}] id={doc_id[:8]}...")
        print(f"    내용: {content[:80]}...")
        print(f"    메타: {meta}\n")
    return count


# 실제 실행: 저장된 컬렉션을 검사
inspect_store(store)
