import gradio as gr
from huggingface_hub import hf_hub_download

# Download model
print("📥 Downloading GGUF model...")
model_path = hf_hub_download(
    repo_id="HexQuant/Pars-Medical-o1-Llama-FFT-GGUF",
    filename="model-Q4_K_M.gguf",  # ✅ Fixed filename
)
print(f"✅ Model downloaded: {model_path}")

from llama_cpp import Llama

print("🔄 Loading model...")
llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2,
    n_gpu_layers=0,
    verbose=False,
)
print("✅ Model ready!")

SYSTEM = """You are Pars-Medical, a bilingual (Persian/English) medical AI assistant with chain-of-thought reasoning.
When answering: analyze carefully, think step-by-step, provide accurate medical information.
Respond in the user's language. ⚠️ Educational purposes only - always consult a real doctor."""


def chat(message, history):
    messages = [{"role": "system", "content": SYSTEM}]
    for user, assistant in history:
        messages.append({"role": "user", "content": user})
        if assistant:
            messages.append({"role": "assistant", "content": assistant})
    messages.append({"role": "user", "content": message})
    
    response = llm.create_chat_completion(messages=messages, max_tokens=512, temperature=0.7, stream=True)
    
    text = ""
    for chunk in response:
        delta = chunk["choices"][0].get("delta", {}).get("content", "")
        text += delta
        yield text


with gr.Blocks(title="🩺 Pars-Medical-o1", theme=gr.themes.Soft()) as demo:
    gr.HTML("""
    <div style="text-align:center; padding:20px;">
        <h1>🩺 Pars-Medical-o1</h1>
        <h3>دستیار پزشکی هوشمند دوزبانه | Bilingual Medical AI Assistant</h3>
        <p><a href="https://huggingface.co/HexQuant/Pars-Medical-o1-Llama-FFT" target="_blank">🤗 Model</a> • <a href="https://twitter.com/xPOURY4" target="_blank">🐦 @xPOURY4</a> • <a href="https://github.com/xPOURY4" target="_blank">💻 GitHub</a></p>
    </div>
    """)
    
    chatbot = gr.Chatbot(height=450, placeholder="سوال پزشکی بپرسید... | Ask a medical question...")
    msg = gr.Textbox(placeholder="پیام خود را بنویسید... | Type your message...", show_label=False)
    
    with gr.Row():
        send = gr.Button("📤 ارسال | Send", variant="primary")
        clear = gr.Button("🗑️ پاک کردن | Clear")
    
    gr.Examples([
        ["علائم دیابت نوع ۲ چیست؟"],
        ["What are symptoms of hypothyroidism?"],
        ["تفاوت میگرن و سردرد تنشی چیست؟"],
        ["How to differentiate viral vs bacterial infection?"],
    ], inputs=msg, label="💡 نمونه سوالات | Examples")
    
    gr.HTML("""
    <div style="background:#fff3cd; padding:12px; border-radius:8px; text-align:center; margin-top:15px;">
        ⚠️ <b>توجه:</b> فقط برای آموزش است، جایگزین مشاوره پزشکی نیست<br>
        <b>Disclaimer:</b> Educational only, not medical advice
    </div>
    <p style="text-align:center; color:#666; margin-top:15px;">Made with ❤️ by <a href="https://twitter.com/xPOURY4">@xPOURY4</a> in Iran 🇮🇷</p>
    """)
    
    msg.submit(chat, [msg, chatbot], chatbot).then(lambda: "", None, msg)
    send.click(chat, [msg, chatbot], chatbot).then(lambda: "", None, msg)
    clear.click(lambda: None, None, chatbot)

demo.queue().launch()