import os
import streamlit as st
from tasks.data_ingestion_task.data_ingestion_task import DataIngestionTask
from tasks.query_handling_task.query_handling_task import QueryHandlingTask
from utils.llama_index_utils import setup_directories

# Title Section
st.title("(PDF) Chat con documentos de Procesos 🗒️")

# Subtitle Section
st.markdown("Retrieval-Augmented Generation")
st.markdown("iniciar chat ...🚀")

# Session State Initialization
if 'messages' not in st.session_state:
    st.session_state.messages = [{'role': 'assistant', "content": 'Hola! Selecciona un pdf para cargar, y hazme una pregunta.'}]

# Sidebar Section
with st.sidebar:
    st.image('image_logo.jpeg', use_column_width=True)
    st.title("Menu:")
    
    uploaded_file = st.file_uploader("Sube un archivo PDF y dale click al botón enviar y procesar.")

    if st.button("Enviar y Procesar"):
        if uploaded_file is not None:
            with st.spinner("Procesando..."):
                # Ensure the data directory exists
                data_dir, persist_dir = setup_directories()
                
                # Save the uploaded file in the data directory
                filepath = os.path.join(data_dir, "saved_pdf.pdf")
                with open(filepath, "wb") as f:
                    f.write(uploaded_file.getbuffer())
                
                # Use DataIngestionTask to process the PDF
                ingestion_task = DataIngestionTask(
                    config_path='tasks/data_ingestion_task/config.txt',
                    input_structure_path='tasks/data_ingestion_task/input_structure.json',
                    output_structure_path='tasks/data_ingestion_task/output_structure.json'
                )
                ingestion_task.execute({})
                
                st.success("PDF processed successfully")
        else:
            st.error("Please upload a PDF file before processing.")

    st.subheader("Select Model:")
    model_selected = st.session_state.get('model_selected', '')
    
    col1, col2 = st.columns(2)
    
    if col1.button("GEMMA", key="gemmabtn"):
        st.session_state.model_selected = "GEMMA"
    
    if col2.button("GEMINI", key="geminibtn"):
        st.session_state.model_selected = "GEMINI"
    
    if model_selected:
        st.write(f"Selected Model: **{model_selected}**")

# Chat Input Section
user_prompt = st.chat_input("Pregunta acerca del contenido en el archivo PDF:")
if user_prompt:
    st.session_state.messages.append({'role': 'user', "content": user_prompt})
    
    # Use QueryHandlingTask to handle the query
    query_task = QueryHandlingTask(
        config_path='tasks/query_handling_task/config.txt',
        input_structure_path='tasks/query_handling_task/input_structure.json',
        output_structure_path='tasks/query_handling_task/output_structure.json'
    )
    response = query_task.execute({'query': user_prompt})
    
    st.session_state.messages.append({'role': 'assistant', "content": response})

# Chat Message Display
for message in st.session_state.messages:
    with st.chat_message(message['role']):
        st.write(message['content'])