import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain_community.embeddings import BedrockEmbeddings
from langchain_aws import ChatBedrock
from langchain_community.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
import os
import zipfile

#el archivo ruta_zip debe existir y es el directorio compreso con la base de datos vectorial Chromadb
def descomprimir_archivo_zip(ruta_zip, ruta_destino):
    """
    Descomprime un archivo ZIP en la ruta de destino especificada.

    :param ruta_zip: Ruta del archivo ZIP a descomprimir.
    :param ruta_destino: Directorio donde se descomprimirá el archivo.
    """
    # Verificar si la ruta de destino existe, si no, crearla
    if not os.path.exists(ruta_destino):
        os.makedirs(ruta_destino)
    
    # Abrir el archivo ZIP
    with zipfile.ZipFile(ruta_zip, 'r') as zip_ref:
        # Extraer todos los archivos en la ruta de destino
        zip_ref.extractall(ruta_destino)
        print(f'Archivos extraídos en: {ruta_destino}')


def cargar_base_datos(directorio_data,embedder):
    """
    Carga una base de datos en el directorio especificado usando Chroma.

    :param directorio_data: Ruta del directorio donde se encuentra la base de datos.
    """
    # Verificar si la ruta del directorio de datos existe
    if not os.path.exists(directorio_data):
        raise FileNotFoundError(f'El directorio {directorio_data} no existe.')

    # Cargar la base de datos utilizando Chroma
    vectorstore = Chroma(persist_directory =directorio_data,embedding_function=embedder)
    
    print(f'Base de datos cargada desde: {directorio_data}')
    return vectorstore





#Inicializar LLM y Embedding de Amazon BedRock
#Las variables de ambiente AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_DEFAULT_REGION se deben configurar 
#en los secretos y variables de ambiente de HuggingFace
def initLLM():
    return ChatBedrock(model_id="anthropic.claude-3-sonnet-20240229-v1:0")
def initEmbedder():
   return BedrockEmbeddings(model_id='amazon.titan-embed-text-v1')
 


# Función para manejar el chat
def handle_chat(message, history):
    # Ejemplo sencillo de respuesta de chat
    print(chromaDB)
    pdf_qa=ConversationalRetrievalChain.from_llm(llm=initLLM(),
                                             retriever=chromaDB.as_retriever(search_kwargs={'k':10}),
                                             verbose=True)
    chat_history = []
    resultado=pdf_qa({"question":message,"chat_history": chat_history})
    return resultado["answer"]


# Definir la interfaz de Gradio
with gr.Blocks() as demo:
        # Elementos para el chat
        chatbot = gr.ChatInterface(fn=handle_chat, undo_btn=None,clear_btn=None,retry_btn=None)
       

   
# Ejecutar la aplicación
if __name__ == "__main__":
    bedrock_llm=initLLM()
    bedrock_embedder=initEmbedder()
    ruta_zip = './data.zip'
    ruta_destino = '.'
    print("Extrayendo base vectorial")
    descomprimir_archivo_zip(ruta_zip, ruta_destino)
    directorio_data = 'data'  # Ruta del directorio de datos
    chromaDB = cargar_base_datos(directorio_data,bedrock_embedder)
    demo.launch()