import gradio as gr
from openai import OpenAI
import re

client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="not-needed")

def strip_think_tags(text):
    return re.sub(r"<think>.*?</think>\s*", "", text, flags=re.DOTALL).strip()

def chat(message, history, system_prompt, max_tokens, temperature, show_thinking):
    messages = []
    if system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})
    for h in history[-4:]:
        if h["role"] == "user":
            messages.append({"role": "user", "content": h["content"]})
        elif h["role"] == "assistant":
            messages.append({"role": "assistant", "content": h["content"]})
    messages.append({"role": "user", "content": message})

    try:
        stream = client.chat.completions.create(
            model="local",
            messages=messages,
            max_tokens=int(max_tokens),
            temperature=float(temperature),
            stream=True,
        )
        partial = ""
        for chunk in stream:
            delta = chunk.choices[0].delta.content or ""
            partial += delta
            if show_thinking:
                yield partial
            else:
                yield strip_think_tags(partial)
    except Exception as e:
        yield f"Error: {str(e)} - llama-server is still loading, please wait and try again."

with gr.Blocks(title="Qwen3.5-9B Claude Reasoning", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Qwen3.5-9B Claude 4.6 Opus Reasoning Distilled v2
        CPU inference ဖြစ်တဲ့အတွက် response 30s-2min ကြာနိုင်ပါတယ်။ မြန်မာလိုလည်း မေးလို့ရပါတယ်။
        """
    )
    system_prompt = gr.Textbox(
        label="System Prompt",
        value="You are a helpful coding and reasoning assistant. Respond in the same language the user uses.",
        lines=2,
    )
    with gr.Row():
        max_tokens = gr.Slider(64, 2048, 512, step=64, label="Max Tokens")
        temperature = gr.Slider(0.1, 1.5, 0.7, step=0.1, label="Temperature")
        show_thinking = gr.Checkbox(label="Show thinking", value=False)

    gr.ChatInterface(
        fn=chat,
        additional_inputs=[system_prompt, max_tokens, temperature, show_thinking],
        examples=[
            ["Python မှာ decorator ဘယ်လိုသုံးရလဲ"],
            ["Write a Flask REST API with JWT auth"],
            ["FizzBuzz ကို recursive function နဲ့ ရေးပြပါ"],
        ],
    )

demo.launch(server_name="0.0.0.0", server_port=7860)