import spaces  # debe importarse antes que cualquier paquete relacionado con CUDA
import gradio as gr
import fitz  # PyMuPDF para leer PDFs
import os

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from groq import Groq

# 1. Configuración de embeddings (Modelo ligero y rápido)
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")


# 2. Función para cargar y procesar el PDF
def load_and_index_pdf(pdf_file):
    if pdf_file is None:
        return "Por favor, sube un archivo PDF."

    # Gradio entrega una ruta de archivo (string) directamente,
    # no hace falta guardar una copia manual.
    doc = fitz.open(pdf_file)
    text = ""
    for page in doc:
        text += page.get_text()

    # Dividir en trozos (Chunking)
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    chunks = text_splitter.split_text(text)

    # Crear la base de datos vectorial FAISS
    vectorstore = FAISS.from_texts(chunks, embeddings)

    return vectorstore


# Función "dummy" solo para satisfacer el requisito de arranque de ZeroGPU
# (Hugging Face exige que exista al menos una función @spaces.GPU, pero
# esta función real de RAG no necesita GPU: FAISS corre en CPU y el LLM
# se llama por API remota, no localmente).
@spaces.GPU
def _zerogpu_startup_check():
    return True


# 3. Función de respuesta (RAG)
def respond(vectorstore, user_query):
    if not user_query:
        return "Hazme una pregunta sobre el documento."

    # Buscar los trozos más relevantes
    docs = vectorstore.similarity_search(user_query, k=2)

    # Construir el contexto
    context = "\n\n".join([d.page_content for d in docs])

    prompt = f"""
    Basado en el siguiente contexto, responde a la pregunta.
    Si la respuesta no está en el contexto, di "No tengo esa información en el documento".

    Contexto: {context}
    Pregunta: {user_query}
    Respuesta:
    """

    groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
    completion = groq_client.chat.completions.create(
        messages=[{"role": "user", "content": prompt}],
        model="llama-3.3-70b-versatile",
        temperature=0.7,
    )
    return completion.choices[0].message.content


# 4. Interfaz Gradio
def main():
    # Estado global para la base de datos
    vectorstore_state = None

    def process_upload(pdf_file):
        nonlocal vectorstore_state
        if pdf_file:
            vectorstore_state = load_and_index_pdf(pdf_file)
            return "Documento procesado con éxito. ¡Pregunta lo que quieras!", gr.update(interactive=True)
        return "Sube un PDF primero.", gr.update(interactive=False)

    def process_query(query):
        if vectorstore_state is None:
            return "Por favor, sube un PDF primero."
        return respond(vectorstore_state, query)

    with gr.Blocks(title="RAG Demo - Dow Jones Style") as demo:
        gr.Markdown("# 📚 Chat con tus Documentos (RAG + FAISS)")
        gr.Markdown("Sube un PDF y pregunta sobre su contenido. ¡Sin alucinaciones!")

        with gr.Row():
            with gr.Column():
                pdf_input = gr.File(label="Sube tu PDF aquí", file_types=[".pdf"])
                upload_btn = gr.Button("Procesar Documento", variant="primary")
                status_msg = gr.Textbox(label="Estado", interactive=False)

            with gr.Column():
                query_input = gr.Textbox(label="Tu pregunta", placeholder="¿De qué trata este documento?")
                submit_btn = gr.Button("Enviar Pregunta", interactive=False)
                output_msg = gr.Textbox(label="Respuesta de la IA")

        # Lógica de interacción
        upload_btn.click(
            fn=process_upload,
            inputs=pdf_input,
            outputs=[status_msg, submit_btn],
        )

        submit_btn.click(
            fn=process_query,
            inputs=query_input,
            outputs=output_msg,
        )

    demo.launch()


if __name__ == "__main__":
    main()