import chromadb
from chromadb.utils import embedding_functions
import json
import os

# Configurações
CHUNKS_FILE = "fluig_chunks.json"
DB_PATH = "fluig_vector_db"

def init_vector_db():
    # Usar um modelo extremamente leve adequado para o Raspberry Pi
    # 'all-MiniLM-L6-v2' é o padrão ouro para performance vs qualidade
    model_name = "all-MiniLM-L6-v2"
    emb_fn = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_name)

    # Inicializa o cliente persistente (SQLite sob o capô)
    client = chromadb.PersistentClient(path=DB_PATH)
    
    # Cria ou obtém a coleção
    collection = client.get_or_create_collection(
        name="fluig_docs",
        embedding_function=emb_fn,
        metadata={"hnsw:space": "cosine"} # Melhor para busca semântica
    )

    if not os.path.exists(CHUNKS_FILE):
        print(f"Erro: {CHUNKS_FILE} não encontrado. Rode o rag_processor.py primeiro.")
        return

    with open(CHUNKS_FILE, "r", encoding="utf-8") as f:
        chunks = json.load(f)

    print(f"Indexando {len(chunks)} chunks no ChromaDB (isso pode levar alguns minutos no Pi)...")
    
    ids = []
    documents = []
    metadatas = []

    for i, chunk in enumerate(chunks):
        ids.append(f"id_{i}")
        documents.append(chunk["content"])
        metadatas.append({
            "title": chunk["metadata"].get("title", ""),
            "source": chunk["metadata"].get("source", ""),
            "path": chunk["metadata"].get("path", "")
        })

    # Adicionar em lotes para não estourar memória
    batch_size = 100
    for i in range(0, len(documents), batch_size):
        collection.add(
            ids=ids[i:i+batch_size],
            documents=documents[i:i+batch_size],
            metadatas=metadatas[i:i+batch_size]
        )
        print(f"Progresso: {i + len(documents[i:i+batch_size])}/{len(documents)}")

    print(f"Sucesso! Banco vetorial criado em {DB_PATH}")

if __name__ == "__main__":
    init_vector_db()