import gradio as gr
import PyPDF2
import os
import json
import vertexai
from vertexai.generative_models import GenerativeModel, Part, SafetySetting

# Configuración global
generation_config = {
    "max_output_tokens": 8192,
    "temperature": 0,
    "top_p": 0.8,
}

safety_settings = [
    SafetySetting(
        category=SafetySetting.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
        threshold=SafetySetting.HarmBlockThreshold.OFF
    ),
    SafetySetting(
        category=SafetySetting.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
        threshold=SafetySetting.HarmBlockThreshold.OFF
    ),
    SafetySetting(
        category=SafetySetting.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
        threshold=SafetySetting.HarmBlockThreshold.OFF
    ),
    SafetySetting(
        category=SafetySetting.HarmCategory.HARM_CATEGORY_HARASSMENT,
        threshold=SafetySetting.HarmBlockThreshold.OFF
    ),
]

def configurar_credenciales(json_path: str):
    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = json_path

def extraer_texto(pdf_path: str) -> str:
    texto_total = ""
    with open(pdf_path, "rb") as f:
        lector = PyPDF2.PdfReader(f)
        for page in lector.pages:
            texto_total += page.extract_text() or ""
    return texto_total

def parsear_con_llm_con_enumeraciones(texto_pdf: str, model: GenerativeModel) -> dict:
    """
    Usa el LLM para extraer las 'Preguntas' y sus 'Respuestas'.
    Reconoce enumeraciones en 'Preguntas' (ej. '1.' o '1)') y en 'RESPUESTAS'.
    Devuelve JSON con la forma:
    {
      "Pregunta 1": "Texto de la respuesta",
      "Pregunta 2": "Texto de la respuesta",
      ...
    }
    """
    prompt = f"""
Eres un parser de texto que recibe el contenido de un PDF con dos secciones:
'Preguntas' y 'RESPUESTAS', cada una enumerada como '1.', '2)', etc.
Tu tarea es emparejar cada pregunta con su respuesta correspondiente, basándote
en el número que las identifica (1, 2, 3, ...).
Devuelve un JSON con la estructura:
{{
  "Pregunta 1": "texto de la respuesta",
  "Pregunta 2": "texto de la respuesta",
  ...
}}
Si no hay correspondencia entre pregunta y respuesta, la dejas vacía.
Ejemplo de JSON vacío: {{}}

Texto PDF:
{texto_pdf}

Devuelve solo el JSON, sin explicaciones adicionales.
"""
    part_text = Part.from_text(prompt)
    response = model.generate_content(
        [part_text],
        generation_config=generation_config,
        safety_settings=safety_settings,
        stream=False
    )
    try:
        data = json.loads(response.text.strip())
        if isinstance(data, dict):
            return data
        else:
            return {}
    except:
        return {}

def comparar_preguntas_respuestas(dict_docente: dict, dict_alumno: dict) -> str:
    """Compara dict_docente vs dict_alumno y retorna retroalimentación."""
    retroalimentacion = []
    for pregunta, resp_correcta in dict_docente.items():
        resp_alumno = dict_alumno.get(pregunta, None)
        if resp_alumno is None:
            retroalimentacion.append(f"**{pregunta}**\nNo fue asignada al alumno.\n")
        else:
            retroalimentacion.append(
                f"**{pregunta}**\n"
                f"Respuesta del alumno: {resp_alumno}\n"
                f"Respuesta correcta: {resp_correcta}\n"
            )
    return "\n".join(retroalimentacion)

def revisar_examen(json_cred, pdf_docente, pdf_alumno):
    """
    Función convertida en generador para mostrar 'progreso' en Gradio:
    en lugar de return final, se hace yield en varios pasos.
    """
    # Paso 1: Configurar credenciales
    yield "Cargando credenciales..."
    try:
        configurar_credenciales(json_cred.name)

        # Paso 2: Inicializar Vertex AI
        yield "Inicializando Vertex AI..."
        vertexai.init(project="deploygpt", location="us-central1")

        # Paso 3: Leer texto PDF
        yield "Extrayendo texto del PDF del docente..."
        texto_docente = extraer_texto(pdf_docente.name)

        yield "Extrayendo texto del PDF del alumno..."
        texto_alumno = extraer_texto(pdf_alumno.name)

        # Paso 4: Invocar modelo
        yield "Parseando preguntas y respuestas (Docente)..."
        model = GenerativeModel(
            "gemini-1.5-pro-001",
            system_instruction=["Eres un parser estricto."]
        )
        dict_docente = parsear_con_llm_con_enumeraciones(texto_docente, model)

        yield "Parseando preguntas y respuestas (Alumno)..."
        dict_alumno = parsear_con_llm_con_enumeraciones(texto_alumno, model)

        # Paso 5: Comparar
        yield "Comparando respuestas..."
        feedback = comparar_preguntas_respuestas(dict_docente, dict_alumno)

        if len(feedback.strip()) < 5:
            yield "No se encontraron preguntas o respuestas válidas."
            return  # terminamos la función

        # Paso 6: Resumen final
        yield "Generando resumen final..."
        summary_prompt = f"""
        Eres un profesor experto de bioquímica. Te muestro la comparación de preguntas y respuestas:
        {feedback}
        Por favor, genera un breve resumen del desempeño del alumno 
        sin inventar preguntas adicionales.
        """
        summary_part = Part.from_text(summary_prompt)
        summary_resp = model.generate_content(
            [summary_part],
            generation_config=generation_config,
            safety_settings=safety_settings,
            stream=False
        )
        final_result = f"{feedback}\n\n**Resumen**\n{summary_resp.text.strip()}"

        yield final_result

    except Exception as e:
        yield f"Error al procesar: {str(e)}"

# Interfaz Gradio
interface = gr.Interface(
    fn=revisar_examen,
    inputs=[
        gr.File(label="Credenciales JSON"),
        gr.File(label="PDF del Docente"),
        gr.File(label="PDF Alumno")
    ],
    # Cuando la función es un generador, definimos outputs="text" (o gr.Textbox()).
    outputs="text",
    title="Revisión de Exámenes con Enumeraciones (Progreso)",
    description=(
        "Sube tus credenciales, el PDF del docente y del alumno. El LLM "
        "detectará enumeraciones (1., 2), etc.) en 'Preguntas' y 'RESPUESTAS' "
        "y mostrará progreso a medida que avanza."
    )
)

interface.launch(debug=True)