# local_process_embeddings.py
from pathlib import Path
import ollama
import faiss
import numpy as np
import os
import glob
import pickle  # For saving documents list


# (Copy your inicializar_rag function here)
# Make sure it's adapted to just return documents and index without Gradio parts
def inicializar_rag():
    """
    Inicializa o modelo de embeddings, documentos e índice FAISS a partir dos arquivos de capítulo.
    Adaptado para retornar apenas documentos e índice FAISS, sem partes do Gradio.
    """

    # Diretório onde os arquivos de capítulo estão localizados
    base_dir = Path(__file__).resolve().parent
    diretorio_rag = base_dir / "RAG"

    print(f">>> Looking for files under: {diretorio_rag}")

    if not diretorio_rag.exists():
        raise FileNotFoundError(
            f"O diretório {diretorio_rag} não existe. Verifique se ele foi copiado corretamente para o container."
        )

    # Use pathlib's glob for more consistent path handling
    arquivos = list(diretorio_rag.glob("[bsde][1-9].txt"))
    caminhos_arquivos = [str(arquivo) for arquivo in arquivos]  # convert pathlib objects to strings

    documentos = []
    # Carrega e processa cada arquivo de capítulo
    for caminho in caminhos_arquivos:
        with open(caminho, "r", encoding="utf-8") as f:
            conteudo = f.read().split("\n\n")
            documentos.extend(conteudo)

    if not documentos:
        raise ValueError("Nenhum documento foi carregado. Verifique os caminhos dos arquivos.")

    print(f"Trabalhando em {len(documentos)} documentos ...")
    # Gera embeddings usando o Ollama
    embeddings = []
    for index, doc in enumerate(documentos):
        if index % 200 == 0:
            print(f"= = = {index} Embeddings... = = = doc: {doc[:20]}")  # added print of the doc

        try:
            response = ollama.embed(model="nomic-embed-text", input=doc)
            embeddings.append(response["embeddings"][0])  # Access the first element of the embeddings list.
        except Exception as e:
            print(f"Error generating embedding for document {index}: {e}")
            # Handle the error, e.g., skip this document or raise an exception
            embeddings.append(None)  # Append None to keep the lists aligned
            continue

    # Remove any None embeddings (documents that failed to embed)
    embeddings_valid = [emb for emb in embeddings if emb is not None]
    documentos_valid = [doc for i, doc in enumerate(documentos) if embeddings[i] is not None]

    if not embeddings_valid:
        raise ValueError("No valid embeddings were generated.")
    print(f"Trabalhando em {len(embeddings_valid)} embeddings...")
    embeddings_np = np.array(embeddings_valid, dtype=np.float32)

    # Cria o índice FAISS
    dimension = embeddings_np.shape[1]
    faiss_index = faiss.IndexFlatL2(dimension)
    faiss_index.add(embeddings_np)

    return documentos_valid, faiss_index


if __name__ == "__main__":
    print("Inicializando RAG localmente...")
    # Ensure the RAG directory is correctly found relative to this script
    # or use an absolute path for local execution.
    # For example, if RAG is in the same dir as this script:
    # base_dir = Path(__file__).resolve().parent
    # diretorio_rag_local = base_dir / "RAG"

    documentos_global, index_global = inicializar_rag()  # Call your function

    print(f"Total de documentos processados: {len(documentos_global)}")
    print(f"Dimensão do índice FAISS: {index_global.d}, Total de vetores: {index_global.ntotal}")

    # Save the documents list
    with open("documentos.pkl", "wb") as f_docs:
        pickle.dump(documentos_global, f_docs)
    print("Lista de documentos salva em documentos.pkl")

    # Save the FAISS index
    faiss.write_index(index_global, "faiss_index.idx")
    print("Índice FAISS salvo em faiss_index.idx")