import os
import gradio as gr
from openai import OpenAI

# -------------------------------------------------
# 🔐 1. Secure API setup
# -------------------------------------------------
# IMPORTANT:
# Do NOT hardcode your token.
# Instead, go to your Space → Settings → Repository secrets
# Add a secret with:
#   Key: HF_TOKEN
#   Value: your Hugging Face token (e.g. hf_xxx...)
# Then this will safely load it at runtime:
client = OpenAI(
    base_url="https://router.huggingface.co/v1",
    api_key=os.getenv("HF_TOKEN"),
)

# -------------------------------------------------
# 🧩 2. Chat function (with streaming)
# -------------------------------------------------
def chat_with_model(message, history, model_name):
    messages = []
    for human, assistant in history:
        messages.append({"role": "user", "content": human})
        if assistant:
            messages.append({"role": "assistant", "content": assistant})
    messages.append({"role": "user", "content": message})

    full_reply = ""
    stream = client.chat.completions.create(
        model=model_name,
        messages=messages,
        stream=True,
    )

    for chunk in stream:
        delta = chunk.choices[0].delta.content
        if delta:
            full_reply += delta
            yield full_reply  # stream live output

# -------------------------------------------------
# 🎨 3. Gradio interface
# -------------------------------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        <h2 style="text-align:center;">🤖 GPT-OSS Chat (via Hugging Face Router)</h2>
        <p style="text-align:center;">
        Stream responses from Fireworks-AI's GPT-OSS models — 7B, 20B, and 120B — through the Hugging Face Inference Router.
        </p>
        """,
        elem_id="title",
    )

    with gr.Row():
        model_selector = gr.Dropdown(
            label="Select Model",
            choices=[
                "openai/gpt-oss-7b:fireworks-ai",
                "openai/gpt-oss-20b:fireworks-ai",
                "openai/gpt-oss-120b:fireworks-ai",
            ],
            value="openai/gpt-oss-120b:fireworks-ai",
            interactive=True,
        )

    gr.ChatInterface(
        fn=lambda msg, hist: chat_with_model(msg, hist, model_selector.value),
        title="GPT-OSS Chat",
        chatbot=gr.Chatbot(height=500, show_label=False),
        textbox=gr.Textbox(placeholder="Ask me anything...", autofocus=True),
        retry_btn="↩ Retry",
        clear_btn="🧹 Clear Chat",
    )

# -------------------------------------------------
# 🚀 4. Launch app
# -------------------------------------------------
if __name__ == "__main__":
    demo.launch()