import os
import gradio as gr
from ollama import Client
import time

OLLAMA_API_KEY = os.environ.get("OLLAMA_API_KEY")
if not OLLAMA_API_KEY:
    raise ValueError("API Key تنظیم نشده است.")

client = Client(
    host="https://ollama.com",
    headers={"Authorization": f"Bearer {OLLAMA_API_KEY}"}
)

AVAILABLE_MODELS = [
    "gpt-oss:20b-cloud",
    "qwen3-coder:480b-cloud"
]

# تابع استریمینگ
def respond(message, history, model_name):
    messages = []
    for h in history:
        messages.append({"role": "user", "content": h[0]})
        messages.append({"role": "assistant", "content": h[1]})
    messages.append({"role": "user", "content": message})

    bot_message = ""
    yield history, ""  # برای ریست کردن Textbox کاربر
    for part in client.chat(model_name, messages=messages, stream=True):
        bot_message += part["message"]["content"]
        yield history + [(message, bot_message)], ""  # بروزرسانی استریم
        time.sleep(0.01)  # کمی تاخیر برای ظاهر شدن استریم روان

# رابط Blocks
with gr.Blocks() as demo:
    gr.Markdown("# 💬 Ollama Cloud Chatbot (Streamed)")
    model_selector = gr.Dropdown(choices=AVAILABLE_MODELS, value=AVAILABLE_MODELS[0], label="انتخاب مدل")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="پیام خود را وارد کنید")
    send_btn = gr.Button("ارسال")

    send_btn.click(respond, inputs=[msg, chatbot, model_selector], outputs=[chatbot, msg])

demo.launch()