from pathlib import Path
import ollama
import faiss
import gradio as gr
import numpy as np
import os
import glob
import pickle

print("Inicializando...")


# Função para carregar e processar o conhecimento
def carregar_recursos_pre_calculados():
    """Carrega documentos e o índice FAISS pré-calculados."""
    print("Carregando documentos e índice FAISS pré-calculados...")
    try:
        with open("documentos.pkl", "rb") as f:
            documentos = pickle.load(f)
        print(f"Documentos carregados: {len(documentos)} itens.")
    except FileNotFoundError:
        print("Erro: documentos.pkl não encontrado.")
        raise
    except Exception as e:
        print(f"Erro ao carregar documentos.pkl: {e}")
        raise

    try:
        index = faiss.read_index("faiss_index.idx")
        print(f"Índice FAISS carregado. Dimensão: {index.d}, Total de vetores: {index.ntotal}")
    except FileNotFoundError:
        print("Erro: faiss_index.idx não encontrado.")
        raise
    except Exception as e:
        print(f"Erro ao carregar faiss_index.idx: {e}")
        raise

    return documentos, index


print("Carregando RAG pré-calculado...")
try:
    documentos_global, index_global = carregar_recursos_pre_calculados()
except Exception as e:
    print(f"Falha crítica ao carregar recursos pré-calculados: {e}")
    # Handle error appropriately, maybe exit or display an error in Gradio
    raise SystemExit(f"Falha ao carregar recursos: {e}")


# Função para buscar contexto relevante
def buscar_contexto(pergunta, documentos, index, k=5):
    """Busca os k documentos mais relevantes para a pergunta."""
    # Gera embedding para a pergunta usando o Ollama
    response = ollama.embed(model="nomic-embed-text", input=pergunta)
    query_embedding = np.array([response["embeddings"][0]], dtype=np.float32)

    # Busca no índice FAISS
    D, I = index.search(query_embedding, k)
    return [documentos[i] for i in I[0]]


# Função para gerar resposta com o Ollama
def gerar_resposta(frase_entrada, documentos, index):
    """Gera uma resposta baseada no contexto recuperado."""
    contexto = buscar_contexto(frase_entrada, documentos, index)
    contexto_str = "\n".join(contexto)

    # Prompt com instruções detalhadas sobre CIF
    prompt = f"""
        1. Identifique o **conceito significativo** da frase "XXXXXX", ou seja, o propósito ou a ideia central que ela expressa, independentemente de ser uma pergunta ou uma afirmação. 
        2. Para determinar a vinculação de um termo específico com a Classificação Internacional de Funcionalidade (CIF), segundo Cieza et al., 2016, siga as diretrizes a seguir: 
        3. Identifique se o termo pertence ao universo da CIF, de acordo com o documento anexado: 
        - Se o termo é mencionado no arquivo anexado, anote isso. 
        - Se o termo não aparece no arquivo anexado, classifique como "Não coberto."
        4. Verifique a relação do termo com os componentes da CIF: 
        - Funções Corporais: Se o termo se relaciona com códigos iniciados pela letra “b”
        - Estruturas Corporais: Se o termo se relaciona com códigos iniciados pela letra “s” 
        - Atividades e Participação: Se o termo se relaciona com códigos iniciados pela letra “d” 
        - Fatores Ambientais: Se o termo se relaciona com códigos iniciados pela letra “e”
        5. Classifique o termo: Se o termo for mencionado na CIF e tiver vínculo com um dos componentes acima, identifique qual categoria ele pertence. 
        Se o termo for mencionado na CIF, mas não tiver vínculo com nenhuma das categorias existentes, classifique como "Não Definido." 

        **RAG**:
        {contexto_str}

        **Frase**:  
        {frase_entrada}

        **Formato da Resposta**:  
        - **Conceito**: [Descrição do conceito significativo]
        - **Código CIF**: [Código CIF + nomeclatura]
        - **Justificativa**: [Explicação baseada no RAG]
    """
    resposta = ollama.generate(model="gemma3:1b", prompt=prompt)
    return resposta["response"]


# Função principal para a interface Gradio
def interface_rag(frase_entrada):
    """Função chamada pela interface Gradio."""
    resposta = gerar_resposta(frase_entrada, documentos, index)
    return resposta


print("Carregando RAG pré-calculado...")
try:
    documentos, index = carregar_recursos_pre_calculados()
except Exception as e:
    print(f"Falha crítica ao carregar recursos pré-calculados: {e}")
    # Handle error appropriately, maybe exit or display an error in Gradio
    raise SystemExit(f"Falha ao carregar recursos: {e}")

# Configuração da interface Gradio
interface = gr.Interface(
    fn=interface_rag,
    inputs=gr.Textbox(label="Digite sua frase", placeholder="Ex.: O que é função puberal?"),
    outputs=gr.Textbox(label="Resposta"),
    title="Sistema RAG para CIF",
    description="Pergunte algo sobre Funções Corporais, Estruturas Corporais, Atividades e Participação ou Fatores Ambientais com base na CIF.",
)

# Iniciar a interface
print("Gerando Interface...")

interface.launch(
    server_name="0.0.0.0",
    server_port=7860,
    show_error=True,
    debug=True,
)  # optional, so you see any startup errors