"""
Este script realiza las siguientes tareas:
1. Carga archivos desde un directorio especificado utilizando la clase `Loader` del módulo `preprocess`.
2. Procesa los archivos mediante limpieza de texto y división en fragmentos.
3. Genera representaciones vectoriales de los textos utilizando `sentence-transformers`.
4. Almacena los vectores en una base de datos Chroma para su posterior recuperación.
5. Inicializa un modelo y ejecuta una interfaz para interactuar con los datos procesados.
Módulos utilizados:
- `preprocess`: Contiene la clase `Loader` para la carga y preprocesamiento de documentos.
- `vdb`: Se asume que gestiona la base de datos vectorial.
- `model_load`: Módulo para cargar el modelo de machine learning.
- `st`: Se asume que proporciona la interfaz de usuario.
Estructura del código:
1. Define el directorio de los archivos a procesar.
2. Carga los archivos y los procesa si el número de archivos es menor a 2.
3. Si hay múltiples archivos, los procesa en un bucle y concatena los fragmentos.
4. Genera embeddings utilizando `sentence-transformers/all-MiniLM-L12-v2`.
5. Almacena los embeddings en ChromaDB y configura un recuperador basado en similitud.
6. Carga el modelo de machine learning.
7. Inicia la interfaz de usuario..
"""
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from tqdm.auto import tqdm
#from chromadb.utils import embedding_functions
from src.preprocess import Loader
from src.vdb import EmbeddingGen
import src.model_load
import streamlit as st

if __name__=="__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)
            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 en base de datos ChromaDB
    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 y ejecución de la interfaz
    src.model_load.load_model()
    print("Lanzando interfaz")
        # Configuración de la página
    st.set_page_config(
        page_title="MathQA - Asistente de Matemáticas",
        page_icon="🧮",
        layout="centered",
        initial_sidebar_state="expanded"
    )
    
    # Paleta de colores neutra
    primary_color = "#010001"
    secondary_color = "#E7E6E7"
    background_color = "#FBFBFA"
    
    # Estilos CSS
    st.markdown(
        f"""
        <style>
        .stApp {{ background-color: {background_color}; }}
        .stTextInput>div>div>input {{
            color: {primary_color};
            background-color: {secondary_color};
            border-radius: 8px;
        }}
        .stButton>button {{
            color: {primary_color};
            background-color: {secondary_color};
            border-radius: 8px;
            transition: all 0.3s;
        }}
        .history-box {{
            border-left: 4px solid {secondary_color};
            padding: 1rem;
            margin: 1rem 0;
            background-color: {secondary_color};
            border-radius: 8px;
        }}
        </style>
        """,
        unsafe_allow_html=True
    )
    
    # Inicializar historial
    if 'history' not in st.session_state:
        st.session_state.history = []
    
    # Variable auxiliar para gestionar el input
    if 'temp_input' not in st.session_state:
        st.session_state.temp_input = ""
    
    # Título de la aplicación
    st.title("🧮 MathQA - Asistente de Matemáticas")
    st.markdown("")
    
    # Widget de entrada con variable auxiliar
    user_input = st.text_input(
        "Escribe tu pregunta matemática aquí:",
        value=st.session_state.temp_input,
        key="user_input",
        placeholder="Ej: ¿Que es una integral?"
    )
    
    # Botón de acción
    col1, col2, col3 = st.columns([5, 4, 4])  # Columnas vacías a los lados para centrar
    with col2:
        if st.button("Resolver pregunta"):
            if user_input:  # Accedemos al valor ingresado
                # Simular respuesta
                mock_answer = src.model_load.ask(user_input, retriever)
                
                # Agregar al historial
                st.session_state.history.insert(0, (user_input, mock_answer))
                
                # Limpiar la variable auxiliar
                st.session_state.temp_input = ""
                
                # Forzar actualización
                st.rerun()
    
    # Mostrar historial
    if st.session_state.history:
        st.markdown("---")
        st.subheader("Historial de Consultas")
        
        for idx, (pregunta, respuesta) in enumerate(st.session_state.history):
            with st.container():
                st.markdown(
                    f"""
                    <div class="history-box">
                        <strong>Pregunta {len(st.session_state.history) - idx}:</strong>
                        <p>{pregunta}</p>
                        <strong>Respuesta:</strong>
                        <p>{respuesta}</p>
                    </div>
                    """,
                    unsafe_allow_html=True
                )
    
    # Pie de página
    st.markdown("---")
    st.markdown("🔍 ¿Necesitas ayuda con álgebra, cálculo o geometría? ¡Estoy aquí para ayudarte!")