import gradio as gr
from huggingface_hub import InferenceClient
import threading
import time

# Hugging Face Model Client
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Thread-local client handling
thread_local = threading.local()

def get_client():
    if not hasattr(thread_local, "client"):
        thread_local.client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
    return thread_local.client

def respond(message, history, system_message, max_tokens, temperature, top_p):
    try:
        time.sleep(0.3)  # Small UX pause
        messages = [{"role": "system", "content": system_message}]
        if history:
            for user, bot in history:
                messages.append({"role": "user", "content": user})
                messages.append({"role": "assistant", "content": bot})
        messages.append({"role": "user", "content": message})

        response = ""
        client = get_client()
        for msg in client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = msg.choices[0].delta.content
            if token:
                response += token
                yield response
    except Exception as e:
        yield f"⚠️ Error: {str(e)}"

def create_interface():
    with gr.Blocks(css="""
        body {
            background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
            font-family: 'Inter', sans-serif;
        }
        .gradio-container {
            max-width: 800px;
            margin: auto;
            padding: 30px;
            border-radius: 16px;
            background: rgba(255, 255, 255, 0.92);
            box-shadow: 0 8px 30px rgba(0, 0, 0, 0.1);
        }
        h1 {
            text-align: center;
            font-size: 2.2em;
            margin-bottom: 0.6em;
        }
        .gr-textbox textarea {
            font-size: 1em;
        }
    """) as demo:

        gr.Markdown("# 🤖 Zephyr Chat Assistant")

        system_message = gr.Textbox(
            value="You are a helpful AI assistant.",
            label="System Instructions",
            lines=2
        )
        max_tokens = gr.Slider(1, 2048, value=512, step=1, label="Maximum Tokens")
        temperature = gr.Slider(0.1, 4.0, value=0.7, step=0.1, label="Temperature")
        top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")

        chatbot = gr.ChatInterface(
            respond,
            additional_inputs=[system_message, max_tokens, temperature, top_p],
            type="messages"
        )

        gr.Examples(
            examples=[
                ["Tell me a short story about a robot learning to paint"],
                ["Explain quantum computing like I'm 10 years old"],
                ["What are some healthy breakfast ideas?"]
            ],
            inputs=[chatbot.textbox]
        )

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(share=True)