import gradio as gr
from transformers import pipeline
from langchain_community.llms import HuggingFacePipeline
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
import os
from huggingface_hub import login

# --- Configurações Iniciais ---
DOC_PATH = "meu-documento.pdf"
MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"

# 🔑 Autenticação no Hugging Face Hub
huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
if huggingface_token:
    try:
        print("Tentando fazer login no Hugging Face Hub...")
        login(token=huggingface_token)
        print("Login bem-sucedido.")
    except Exception as e:
        print(f"Erro no login do Hugging Face: {e}")
else:
    print("Aviso: Variável de ambiente HUGGINGFACE_TOKEN não encontrada. O carregamento do modelo pode falhar.")

# 🔁 Carrega e processa o PDF
def carregar_documento_pdf(path):
    if not os.path.exists(path):
        print(f"Erro: O arquivo {path} não foi encontrado.")
        return None
    loader = PyMuPDFLoader(path)
    documentos = loader.load()
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=256,
        chunk_overlap=100,
        separators=["\n\n", "\n", ".", " "]
    )
    documentos_divididos = splitter.split_documents(documentos)
    embeddings = HuggingFaceEmbeddings(model_name="intfloat/e5-base-v2")
    vectorstore = FAISS.from_documents(documentos_divididos, embeddings)
    return vectorstore

# ⚙️ Configura o pipeline com o modelo Mistral
try:
    pipe = pipeline(
        "text-generation",
        model=MODEL_NAME,
        max_new_tokens=256,
        temperature=0.7,
        do_sample=True
    )
    llm = HuggingFacePipeline(pipeline=pipe)
except Exception as e:
    print(f"Erro ao carregar o modelo HuggingFace: {e}")
    llm = None
    
# 🧠 Memória da conversação
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True, k=3)

# 📄 Indexação do documento
vectorstore = carregar_documento_pdf(DOC_PATH)

# 🧾 Prompt otimizado para respostas claras
custom_prompt = PromptTemplate.from_template("""
    Você é um assistente técnico que responde com base **apenas** no documento fornecido.
    Sempre use o idioma Português do Brasil para responder.

    Contexto:{context}
    Pergunta: {question}
    Resposta:""")

# 🔗 Cadeia RAG com memória e prompt
chat_chain = None
if llm and vectorstore:
    chat_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=vectorstore.as_retriever(search_kwargs={"k": 4}),
        memory=memory,
        combine_docs_chain_kwargs={"prompt": custom_prompt}
    )

# 💬 Interface Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 🤖💼 CS Search Engine")
    chatbot = gr.Chatbot(label="Assistente Técnico")
    
    with gr.Row():
        entrada = gr.Textbox(
            placeholder="Pergunte sobre os projetos do documento...",
            show_label=False,
            container=False,
        )
        enviar_btn = gr.Button("Enviar")

    def interagir(msg, history):
        if not chat_chain:
            history.append((msg, "Erro: A cadeia de chat não foi inicializada."))
            return "", history

        try:
            resposta_completa = chat_chain.invoke({"question": msg})
            resposta_texto_bruto = resposta_completa['answer']
            
            # --- Correção para remover o prompt duplicado (Opção 1) ---
            # Procura a string "Resposta:" para extrair apenas a parte desejada
            partes = resposta_texto_bruto.split("Resposta:")
            if len(partes) > 1:
                # Pega a parte do texto após "Resposta:" e remove espaços em branco extras
                resposta_texto = partes[1].strip()
            else:
                # Se não encontrar "Resposta:", usa a resposta bruta
                resposta_texto = resposta_texto_bruto
            # --- Fim da correção ---
            
        except Exception as e:
            resposta_texto = f"Erro ao gerar resposta: {e}"
        
        history.append((msg, resposta_texto))
        
        return "", history

    entrada.submit(interagir, [entrada, chatbot], [entrada, chatbot])
    enviar_btn.click(interagir, [entrada, chatbot], [entrada, chatbot])

demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)