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

@spaces.GPU
def keep_zero_alive():
    return "ZERO is alive"

client = InferenceClient(
    model="zai-org/GLM-5.3-Flash",
    api_key=os.environ.get("HF_TOKEN")
)

def chat_fn(message, history):
    keep_zero_alive()

    messages = []
    # Last 3 chat hi rakhenge
    if len(history) > 3:
        history = history[-3:]

    for h in history:
        # Purana gradio tuple deta hai: [user, assistant]
        if isinstance(h, (list, tuple)):
            if h[0]:
                messages.append({"role": "user", "content": h[0]})
            if len(h) > 1 and h[1]:
                messages.append({"role": "assistant", "content": h[1]})
        # Naya gradio dict deta hai
        elif isinstance(h, dict):
            messages.append(h)

    messages.append({"role": "user", "content": message})

    try:
        res = client.chat_completion(
            messages=messages,
            max_tokens=2048,
            temperature=0.8
        )
        return res.choices[0].message.content
    except Exception as e:
        return f"⚠️ Server busy, 2 sec baad try karo. Error: {str(e)[:150]}"

with gr.Blocks() as demo:
    gr.Markdown("# GLM-5.3-Flash ⚡ 321B - First on HF!")
    # type hata diya, ab error nahi aayega
    gr.ChatInterface(fn=chat_fn, title="GLM-5.3-Flash ⚡ 321B - First on HF!")

demo.launch()