import re
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

print("Downloading model...")
model_path = hf_hub_download(
    repo_id="Jackrong/Qwen3.5-9B-Claude-4.6-Opus-Reasoning-Distilled-v2-GGUF",
    filename="Qwen3.5-9B.Q4_K_M.gguf",
)
print(f"Model: {model_path}")

print("Loading model...")
llm = Llama(
    model_path=model_path,
    n_ctx=2048,
    n_threads=2,
    n_batch=256,
    verbose=False,
)
print("Model loaded!")

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:
        response = llm.create_chat_completion(
            messages=messages,
            max_tokens=int(max_tokens),
            temperature=float(temperature),
            stream=True,
        )
        partial = ""
        for chunk in response:
            delta = chunk["choices"][0]["delta"].get("content", "")
            partial += delta
            if show_thinking:
                yield partial
            else:
                yield strip_think_tags(partial)
    except Exception as e:
        yield f"Error: {str(e)}"

with gr.Blocks(title="Qwen3.5-9B Claude Reasoning") as demo:
    gr.Markdown("# Qwen3.5-9B Claude 4.6 Opus Reasoning v2\nCPU inference - response ကြာနိုင်ပါတယ်")
    system_prompt = gr.Textbox(label="System Prompt", value="You are a helpful 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])

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