from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from tqdm.auto import tqdm
from src.preprocess import Loader
from src.vdb import EmbeddingGen
import src.model_load
import gradio as gr
import spaces

#@spaces.GPU  # Ubica aquí la anotación correctamente
def main():
    # Definición de directorio
    archivo = [r"data/Calculo_Trascendentes_Tempranas_Zill_4t.pdf"]
    
    # Carga de archivos y procesamiento de texto
    if len(archivo) < 2:
        Load = Loader(archivo[0])
        documentos = Load.load_docs()
        textos_limpios = [Loader.limpiar_texto(texto=doc) for doc in documentos]
        textos = Loader.splitter(texto=textos_limpios, chunk_size=500, chunk_overlap=50)
    else:
        textos = []
        for i in range(len(archivo)):
            Load = Loader(archivo[i])
            documentos = Load.load_docs()
            textos_limpios = [Load.limpiar_texto(texto=doc) for doc in documentos]
            chunks = Load.splitter(texto=textos_limpios, chunk_size=500, chunk_overlap=50)
            textos.extend(chunks)
    
    # Generación de embeddings y almacenamiento
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L12-v2")
    persist_directory = "./persist_directory"
    db = Chroma(persist_directory=persist_directory, embedding_function=embeddings.embed_documents)
    vectorstore = Chroma.from_documents(textos[:10], embeddings, persist_directory="./chroma_db")    
    print("Vectorizado terminado")
    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
    print("Carga del modelo")
    
    # Carga del modelo
    qa_chain = src.model_load.load_model()
    print("Lanzando interfaz")
    
    def chatbot_response(message):
        return qa_chain(message, retriever)

    iface = gr.Interface(
        fn=chatbot_response,
        inputs=gr.Textbox(label="Escribe tu mensaje"),
        outputs=gr.Textbox(label="Respuesta del chatbot"),
        title="Chatbot Simple",
        description="Interfaz simple de chatbot con Gradio."
    )
    
    iface.launch(server_port=7860, share=True)  # Aquí puedes asegurarte de que `share=True` esté activo


if __name__ == "__main__":
    main()