import gradio as gr
from huggingface_hub import InferenceClient


def respond(
    message,
    history: list[dict[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    hf_token: gr.OAuthToken,
):
    """
    Myanmar Coding Assistant - amk-coder-v2
    Powered by HuggingFace Inference API
    """
    client = InferenceClient(token=hf_token.token, model="amkyawdev/amk-coder-v2")

    messages = [{"role": "system", "content": system_message}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        choices = message.choices
        token = ""
        if len(choices) and choices[0].delta.content:
            token = choices[0].delta.content
        response += token
        yield response


# Enhanced Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("""
    # 🤖 amk-coder-v2
    ### Myanmar Coding Assistant | Powered by Qwen2.5-Coder-1.5B
    
    🇲🇲 မြန်မာဘာသာစကားနှင့် English နှစ်ခုလုံးတွင် code ရေးပေးပါသည်
    """)
    
    gr.Markdown("""
    ### ✨ Features
    - 🇲🇲 **Myanmar Support** - Full Myanmar Unicode text support
    - 💻 **Code Generation** - Python, JavaScript, C++, Java, and more
    - 🐛 **Debugging** - Bug detection and fixes
    - 📖 **Code Explanation** - Line-by-line explanations
    """)
    
    gr.Markdown("""
    ### 💡 Example Prompts
    - 🇲🇲 `Python function တစ်ခုရေးပါ။ list sorting လုပ်ပေးပါ။`
    - 🇲🇲 `JavaScript နဲ့ API call ရေးပါ`
    - 🇬🇧 `Write a Python function to reverse a string`
    - 🐛 `Fix this: print('Hello' + 5)`
    """)
    
    with gr.Sidebar():
        gr.LoginButton()
        gr.Markdown("""
        ### ⚙️ Settings
        Adjust generation parameters below
        """)
    
    gr.ChatInterface(
        respond,
        additional_inputs=[
            gr.Textbox(
                value="""You are amk-coder-v2, a helpful Myanmar coding assistant.
- Respond in the same language as the user (Myanmar or English)
- Provide clean, well-commented code
- Explain code when helpful""",
                label="System Prompt",
                lines=4
            ),
            gr.Slider(
                minimum=64,
                maximum=2048,
                value=512,
                step=64,
                label="Max New Tokens"
            ),
            gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.2,
                step=0.1,
                label="Temperature (lower = more focused)"
            ),
            gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-p (nucleus sampling)"
            ),
        ],
        chatbot=gr.Chatbot(height=450),
        title="",
        description="",
    )

    gr.Markdown("""
    ---
    **amk-coder-v2** | Fine-tuned from Qwen2.5-Coder-1.5B
    🇲🇲 Made with ❤️ for Myanmar Developers
    [📦 Model Card](https://huggingface.co/amkyawdev/amk-coder-v2) · Apache-2.0 License
    """)


if __name__ == "__main__":
    demo.launch()