import os
import gradio as gr
import requests
from huggingface_hub import InferenceClient

"""
For more information on `huggingface_hub` Inference API support, 
please check the docs: 
https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""

# ----------------------------------------------------------------
# CONFIGURACIÓN DE SERPER (búsqueda web)
# ----------------------------------------------------------------
SERPER_API_KEY = os.getenv("SERPER_API_KEY")

def do_websearch(query: str) -> str:
    """ Llama a serper.dev para hacer la búsqueda en Google y devuelve texto. """
    if not SERPER_API_KEY:
        return "(SERPER_API_KEY no está configurado)"

    url = "https://google.serper.dev/search"
    headers = {"X-API-KEY": SERPER_API_KEY, "Content-Type": "application/json"}
    payload = {"q": query}

    try:
        resp = requests.post(url, json=payload, headers=headers, timeout=10)
        data = resp.json()
    except Exception as e:
        return f"(Error al llamar a serper.dev: {e})"

    if "organic" not in data:
        return "No se encontraron resultados en serper.dev."

    results = data["organic"]
    if not results:
        return "No hay resultados relevantes."

    text = []
    for i, item in enumerate(results, start=1):
        title = item.get("title", "Sin título")
        link = item.get("link", "Sin enlace")
        text.append(f"{i}. {title}\n   {link}")

    return "\n".join(text)

# ----------------------------------------------------------------
# CONFIGURACIÓN DEL MODELO (AHORA CON TOKEN)
# ----------------------------------------------------------------
hf_api_token = os.getenv("HF_API_TOKEN")
client = InferenceClient(
    model="meta-llama/Llama-3.1-8B-Instruct",
    token=hf_api_token
)

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    use_search  # <-- Checkbox en 1er lugar en additional_inputs
):
    """
    - system_message: Texto del rol "system"
    - history: lista de (user_msg, assistant_msg)
    - message: Mensaje actual del usuario
    - use_search: booleano que indica si se habilita la búsqueda en serper
    """

    # Si use_search es True, primero el contenido web y luego el input
    if use_search:
        web_info = do_websearch(message)
        message = f"Información de la web:\n{web_info}\n\nPregunta del usuario:\n{message}"

    # Construimos la conversación
    messages = [{"role": "system", "content": system_message}]
    for user_txt, assistant_txt in history:
        if user_txt:
            messages.append({"role": "user", "content": user_txt})
        if assistant_txt:
            messages.append({"role": "assistant", "content": assistant_txt})

    # Añadimos turno nuevo
    messages.append({"role": "user", "content": message})

    # Llamamos a la API con streaming
    response_text = ""
    for chunk in client.chat_completion(
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=True
    ):
        token = chunk.choices[0].delta.get("content", "")
        response_text += token
        yield response_text

# ----------------------------------------------------------------
# CONFIGURACIÓN DE LA INTERFAZ
# ----------------------------------------------------------------
# Subimos la casilla de verificación (checkbox) de websearch a la parte superior
demo = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        # 1) Checkbox de websearch
        gr.Checkbox(
            value=False,
            label="🌐 Búsqueda",
        ),
        # 2) Texto system
        gr.Textbox(
            value=(
                "Eres Juan, un asistente virtual en español. "
                "Debes responder con mucha paciencia y empatía a usuarios que "
                "pueden tener dificultades cognitivas o escribir frases confusas. "
                "Provee explicaciones simples, procura entender la intención del usuario "
                "aunque la frase esté mal escrita, y mantén siempre un tono amable."
            ),
            label="Mensaje del sistema",
        ),
        # 3) Sliders
        gr.Slider(
            minimum=1, 
            maximum=2048, 
            value=512, 
            step=1, 
            label="Máxima cantidad de tokens"
        ),
        gr.Slider(
            minimum=0.1, 
            maximum=4.0, 
            value=0.7, 
            step=0.1, 
            label="Temperatura"
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (muestreo por núcleo)",
        ),
    ],
)

if __name__ == "__main__":
    demo.launch()