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

# Engaña al verificador de ZeroGPU al arrancar.
# Al no llamarse nunca durante el chat, NO consume ni 1 segundo de tu cuota.
@spaces.GPU
def _dummy_gpu_checker():
    pass

DEFAULT_MODEL_ID = "apetersson/DeepSeek-V4-Flash-0731-Abliterated-FP8"
MODEL_ID = os.environ.get("MODEL_ID", DEFAULT_MODEL_ID)
HF_TOKEN = os.environ.get("HF_TOKEN")

# Se pasa el token como api_key
client = InferenceClient(model=MODEL_ID, api_key=HF_TOKEN)

# SIN @spaces.GPU -> Se ejecuta en la CPU del contenedor de forma ilimitada
def responder(mensaje, historial, temperatura, top_p, max_tokens):
    mensajes = []
    
    for turno in historial:
        if isinstance(turno, dict):
            mensajes.append(turno)
        elif isinstance(turno, (list, tuple)) and len(turno) == 2:
            turno_usuario, turno_asistente = turno
            if turno_usuario:
                mensajes.append({"role": "user", "content": turno_usuario})
            if turno_asistente:
                mensajes.append({"role": "assistant", "content": turno_asistente})
                
    mensajes.append({"role": "user", "content": mensaje})

    respuesta_parcial = ""
    try:
        stream = client.chat_completion(
            messages=mensajes,
            max_tokens=max_tokens,
            temperature=temperatura,
            top_p=top_p,
            stream=True,
        )
        for fragmento in stream:
            delta = fragmento.choices[0].delta.content or ""
            respuesta_parcial += delta
            yield respuesta_parcial
    except Exception as error:
        yield f"⚠️ Error llamando al modelo `{MODEL_ID}`: {error}"


with gr.Blocks(title="DeepSeek V4 Flash — Demo") as demo:
    gr.Markdown(f"## DeepSeek V4 Flash — Demo\nModelo activo: `{MODEL_ID}`")

    with gr.Accordion("Parámetros de generación", open=False):
        temperatura = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label="Temperature")
        top_p = gr.Slider(0.0, 1.0, value=0.95, step=0.01, label="Top-p")
        max_tokens = gr.Slider(64, 4096, value=512, step=64, label="Max tokens")

    gr.ChatInterface(
        fn=responder,
        additional_inputs=[temperatura, top_p, max_tokens],
        title=None,
        type="messages",
    )

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