import streamlit as st
import fitz  # PyMuPDF
from pdf_processor import read_pdf, process_chunks
from embedding import (
    embed_text_chunks,
    save_to_chroma_db,
    query_chroma_db,
    generate_document_id,
)
from llm_query import query_llm
from chroma_setup import initialize_client
import random
import uuid
from streamlit_pdf_viewer import pdf_viewer  # Visor de PDF integrado en Streamlit

# Inicializa el cliente de ChromaDB (base de datos vectorial)
client = initialize_client()

# Frases sarcásticas para mostrar al usuario mientras procesa
sarcastic_lines = [
    "Cocinando el PDF... 🍳",
    "Esto puede tomar un momento; estoy convenciendo al PDF de que coopere... 🤔",
    "Transformando mágicamente el papel en datos... ✨",
    "Agitando mi varita mágica... 🪄",
    "Invocando a los espíritus del PDF... 👻",
    "Un momento, mientras enseño modales a este PDF... 📚",
    "El PDF está contemplando su existencia... 🧘‍♂️",
    "El PDF se resiste, pero lo domesticaré... ¡en breve!",
]

# Asegura un identificador de usuario único en la sesión
if "user_id" not in st.session_state:
    st.session_state.user_id = str(uuid.uuid4())

user_id = st.session_state.user_id

# Historial de conversación
if "qa_history" not in st.session_state:
    st.session_state.qa_history = []

# Configura la página
st.set_page_config(page_title="Asistente PDF", page_icon="📄")

# Título principal
st.title("📄 Asistente PDF: Procesemos tu Documento")
st.markdown("Sube tu PDF en formato de texto y haz preguntas sobre su contenido.")

# Paso 1: Subir archivo
uploaded_file = st.file_uploader("Subir PDF", type="pdf")

if uploaded_file:
    # Genera un ID único de documento
    document_id = generate_document_id()
    st.success("¡Archivo subido con éxito! Podemos comenzar.")

    # Muestra una frase sarcástica
    st.markdown(random.choice(sarcastic_lines))

    pdf_bytes = uploaded_file.read()  # Lee los bytes del archivo subido

    # Muestra el PDF usando el visor
    st.subheader("PDF Subido")
    pdf_viewer(pdf_bytes)

    # Paso 2: Procesar PDF (extraer texto)
    with st.spinner("Procesando..."):
        st.markdown(random.choice(sarcastic_lines))
        pages_and_text = read_pdf(pdf_bytes)
    st.success("¡PDF procesado con éxito!")

    # Paso 3: Dividir texto en oraciones y chunks
    with st.spinner("Dividiendo en secciones..."):
        st.markdown(random.choice(sarcastic_lines))
        processed_chunks = process_chunks(pages_and_text)
    st.success("¡Texto dividido en chunks!")

    # Paso 4: Generar embeddings localmente
    with st.spinner("Creando embeddings..."):
        st.markdown(random.choice(sarcastic_lines))
        embeddings_df = embed_text_chunks(processed_chunks)
    st.success("¡Embeddings generados!")

    # Paso 5: Guardar embeddings en ChromaDB
    with st.spinner("Guardando embeddings..."):
        st.markdown(random.choice(sarcastic_lines))
        save_to_chroma_db(embeddings_df, user_id, document_id)
    st.success(f"¡Embeddings guardados para el documento con ID: {document_id}")

    # Paso 6: Interacción con el chat
    # Campo de texto para que el usuario haga su pregunta
    query = st.chat_input("Pregunta algo sobre el documento")

    if query:
        with st.spinner("Buscando respuestas..."):
            # Opcional: si quieres ver cuál es el contexto extraído, puedes llamar a:
            # context = query_chroma_db(user_id, document_id, query)
            response = query_llm(query, user_id, document_id)
            # Agrega la interacción al historial
            st.session_state.qa_history.append({"question": query, "answer": response})

# Mostrar historial de chat
if st.session_state.qa_history:
    st.subheader("Historial de Conversación")
    for interaction in st.session_state.qa_history:
        with st.chat_message("user"):
            st.write(interaction["question"])
        with st.chat_message("assistant"):
            st.write(interaction["answer"])

# Botón para borrar el historial
if st.button("Borrar Historial"):
    st.session_state.qa_history = []
    st.success("¡Historial de conversación borrado!")