"""Setup ChromaDB with metadata from metadata.jsonl"""
import json
import os
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.schema import Document

def setup_vector_store():
    """Load metadata.jsonl and populate ChromaDB"""
    print("Loading embeddings model...")

    # Set HuggingFace mirror for China (optional)
    # os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'

    # Use a smaller, faster embedding model that's less likely to timeout
    embeddings = HuggingFaceEmbeddings(
        model_name="all-MiniLM-L6-v2",  # Smaller model, faster download
        model_kwargs={'device': 'cpu'},
        encode_kwargs={'normalize_embeddings': True}
    )

    print("Loading metadata.jsonl...")
    documents = []
    with open("metadata.jsonl", "r", encoding="utf-8") as f:
        for i, line in enumerate(f):
            if i >= 50:  # Limit to first 50 for quick setup
                break
            record = json.loads(line)
            # Combine question and answer as the document content
            content = f"Question: {record.get('Question', '')}\nFinal Answer: {record.get('Final answer', '')}"
            doc = Document(
                page_content=content,
                metadata={
                    "task_id": record.get("task_id", ""),
                    "question": record.get("Question", "")
                }
            )
            documents.append(doc)

    print(f"Creating ChromaDB with {len(documents)} documents...")
    vector_store = Chroma.from_documents(
        documents=documents,
        embedding=embeddings,
        collection_name="gaia_questions",
        persist_directory="./chroma_db"
    )

    print("✅ ChromaDB setup complete!")
    print(f"Stored {len(documents)} question-answer pairs")

    # Test retrieval
    print("\n🔍 Testing retrieval...")
    test_query = "What is the capital of France?"
    results = vector_store.similarity_search(test_query, k=1)
    if results:
        print(f"Query: {test_query}")
        print(f"Similar result: {results[0].page_content[:200]}...")

    return vector_store

if __name__ == "__main__":
    setup_vector_store()