modules/chat_tab.py

import gradio as gr
import fitz  # pymupdf
import requests
import io
import asyncio
import os
from backend.synthesis import SynthesisEngine, PROVIDERS

def extract_pdf_text(url: str) -> str:
    try:
        response = requests.get(url, timeout=15)
        response.raise_for_status()
        doc = fitz.open(stream=response.content, filetype="pdf")
        text = ""
        for page in doc:
            text += page.get_text()
        return text[:20000] # Limit to avoid excessive tokens
    except Exception as e:
        return f"[Error extrayendo PDF: {str(e)}]"

async def handle_classic_chat(message, history, report_md, provider, model):
    if not message.strip():
        yield history
        return
    
    env_key = PROVIDERS.get(provider, {}).get("env_key", "")
    api_key = os.environ.get(env_key, "")
    
    engine = SynthesisEngine(provider=provider, model=model, api_key=api_key)
    
    system_prompt = f"""Eres un asistente de investigación académica.
El usuario ha generado un reporte de investigación.
Responde a las preguntas del usuario BASÁNDOTE ÚNICAMENTE en este reporte.
Si la respuesta no está en el reporte, dilo claramente.

REPORTE ACTUAL:
{report_md}
"""

    history.append([message, "Analizando reporte..."])
    yield history
    
    try:
        history_text = "\n".join([f"User: {h[0]}\nAssistant: {h[1]}" for h in history[:-1]])
        user_prompt = f"Historial:\n{history_text}\n\nNueva pregunta: {message}"
        
        response = await engine._call_llm(system_prompt, user_prompt, temperature=0.0)
        history[-1][1] = response
        yield history
    except Exception as e:
        history[-1][1] = f"❌ Error: {str(e)}"
        yield history

async def handle_paper_chat(message, history, selected_paper, use_pdf, docs_df, provider, model):
    if not message.strip() or docs_df is None or docs_df.empty or not selected_paper:
        yield history
        return
        
    try:
        idx = int(selected_paper.split("-")[0].strip()) - 1
        row = docs_df.iloc[idx]
    except:
        history.append([message, "❌ Error: Selecciona un paper válido de la lista."])
        yield history
        return

    env_key = PROVIDERS.get(provider, {}).get("env_key", "")
    api_key = os.environ.get(env_key, "")
    engine = SynthesisEngine(provider=provider, model=model, api_key=api_key)
    
    context = f"TÍTULO: {row.get('Título', '')}\nAUTORES: {row.get('Autores', '')}\nAÑO: {row.get('Año', '')}\nDOI: {row.get('DOI', '')}\n"
    
    history.append([message, "Preparando contexto..."])
    yield history

    if use_pdf and row.get('PDF URL'):
        history[-1][1] = "Descargando y extrayendo PDF..."
        yield history
        # Ejecutar extracción en un thread para no bloquear el loop asíncrono
        pdf_text = await asyncio.to_thread(extract_pdf_text, row.get('PDF URL'))
        context += f"\nCONTENIDO PDF:\n{pdf_text}"
    else:
        context += f"\nFUENTE: {row.get('Fuente', '')}\nGRADE: {row.get('GRADE', '')}\n"
        if row.get('Abstract'):
            context += f"\nABSTRACT: {row.get('Abstract', '')}"

    system_prompt = f"""Eres un experto analizando papers académicos.
Responde preguntas específicamente sobre este paper basándote en el contexto proporcionado.
Si te preguntan algo que no está en el contexto, indícalo.

CONTEXTO DEL PAPER:
{context}
"""

    history[-1][1] = "Analizando con IA..."
    yield history
    
    try:
        history_text = "\n".join([f"User: {h[0]}\nAssistant: {h[1]}" for h in history[:-1]])
        user_prompt = f"Historial:\n{history_text}\n\nNueva pregunta: {message}"
        
        response = await engine._call_llm(system_prompt, user_prompt, temperature=0.0)
        history[-1][1] = response
        yield history
    except Exception as e:
        history[-1][1] = f"❌ Error: {str(e)}"
        yield history

def update_paper_choices(docs_df):
    if docs_df is None or docs_df.empty:
        return gr.update(choices=[], value=None)
    choices = [f"{i+1} - {row.get('Título', 'Sin título')[:80]}..." for i, row in docs_df.iterrows()]
    return gr.update(choices=choices, value=choices[0] if choices else None)

def create_chat_tabs(report_md_state, docs_df_state, provider_state, model_state):
    """
    Agrega los tabs de chat clásico e IA.
    Retorna None, se asume que se llama dentro de un bloque gr.Tabs()
    """
    with gr.TabItem("💬 Chat"):
        gr.Markdown("### 💬 Chat con el Reporte\nHaz preguntas sobre el informe de investigación generado.")
        chatbot_classic = gr.Chatbot(height=450, elem_classes=["glass-panel"])
        with gr.Row():
            msg_classic = gr.Textbox(placeholder="Pregunta sobre los hallazgos...", show_label=False, scale=4, container=False)
            btn_classic = gr.Button("Enviar", variant="primary", scale=1)
            clear_classic = gr.Button("🗑️ Limpiar", scale=1)
        
        # Events
        msg_classic.submit(
            handle_classic_chat, 
            inputs=[msg_classic, chatbot_classic, report_md_state, provider_state, model_state], 
            outputs=[chatbot_classic]
        ).then(lambda: "", None, [msg_classic])
        
        btn_classic.click(
            handle_classic_chat, 
            inputs=[msg_classic, chatbot_classic, report_md_state, provider_state, model_state], 
            outputs=[chatbot_classic]
        ).then(lambda: "", None, [msg_classic])
        
        clear_classic.click(lambda: [], None, chatbot_classic, queue=False)

    with gr.TabItem("🤖 Chat IA (Paper)"):
        gr.Markdown("### 🤖 Chat con Paper Individual\nSelecciona un paper para analizarlo en profundidad.")
        with gr.Row():
            paper_dropdown = gr.Dropdown(choices=[], label="Seleccionar Paper", scale=3, interactive=True)
            use_pdf = gr.Checkbox(label="📄 Leer PDF (si está disponible)", value=False, scale=1)
            refresh_btn = gr.Button("🔄 Refrescar", scale=1)
            
        chatbot_paper = gr.Chatbot(height=400, elem_classes=["glass-panel"])
        with gr.Row():
            msg_paper = gr.Textbox(placeholder="Pregunta sobre este paper específico...", show_label=False, scale=4, container=False)
            btn_paper = gr.Button("Enviar", variant="primary", scale=1)
            clear_paper = gr.Button("🗑️", scale=1)

        # Events
        refresh_btn.click(update_paper_choices, inputs=[docs_df_state], outputs=[paper_dropdown])
        
        msg_paper.submit(
            handle_paper_chat,
            inputs=[msg_paper, chatbot_paper, paper_dropdown, use_pdf, docs_df_state, provider_state, model_state],
            outputs=[chatbot_paper]
        ).then(lambda: "", None, [msg_paper])
        
        btn_paper.click(
            handle_paper_chat,
            inputs=[msg_paper, chatbot_paper, paper_dropdown, use_pdf, docs_df_state, provider_state, model_state],
            outputs=[chatbot_paper]
        ).then(lambda: "", None, [msg_paper])
        
        clear_paper.click(lambda: [], None, chatbot_paper, queue=False)
        
    return refresh_btn, paper_dropdown