"""
Qwen3 Chat — Hugging Face Space
Deploy by pushing this folder to a Gradio Space on huggingface.co
"""

import os
import torch
import gradio as gr
from threading import Thread
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TextIteratorStreamer,
    BitsAndBytesConfig,
)

# ── Config ────────────────────────────────────────────────────────────────────
MODEL_ID      = os.getenv("MODEL_ID", "Qwen/Qwen3-4B")
MAX_NEW_TOK   = int(os.getenv("MAX_NEW_TOKENS", "512"))
THINKING_MODE = os.getenv("THINKING_MODE", "false").lower() == "true"
USE_4BIT      = os.getenv("USE_4BIT", "false").lower() == "true"
# ──────────────────────────────────────────────────────────────────────────────

print(f"Loading {MODEL_ID}  |  4-bit={USE_4BIT}  |  thinking={THINKING_MODE}")

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

load_kwargs = dict(device_map="auto")

if USE_4BIT:
    load_kwargs["quantization_config"] = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.float16,
    )
    load_kwargs["torch_dtype"] = torch.float16
elif torch.cuda.is_available():
    load_kwargs["torch_dtype"] = torch.float16
else:
    load_kwargs["torch_dtype"] = torch.float32

model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **load_kwargs)
model.eval()
print("Model ready.")


def build_prompt(history: list[dict]) -> str:
    """Convert Gradio-style history to a tokenized prompt string."""
    messages = [{"role": m["role"], "content": m["content"]} for m in history]
    return tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=THINKING_MODE,
    )


def respond(message, history, system_prompt, max_tokens, temperature, top_p):
    # Build message list
    msgs = []
    if system_prompt.strip():
        msgs.append({"role": "system", "content": system_prompt.strip()})
    for h in history:
        msgs.append({"role": "user",      "content": h[0]})
        msgs.append({"role": "assistant", "content": h[1]})
    msgs.append({"role": "user", "content": message})

    prompt = build_prompt(msgs)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )

    gen_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=max_tokens,
        do_sample=temperature > 0,
        temperature=max(temperature, 1e-6),
        top_p=top_p,
        pad_token_id=tokenizer.eos_token_id,
    )

    thread = Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

    partial = ""
    for chunk in streamer:
        partial += chunk
        yield partial


# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(title=f"Qwen3 Chat — {MODEL_ID}", theme=gr.themes.Soft()) as demo:
    gr.Markdown(f"## 🤖 {MODEL_ID}\nStreaming chat powered by Qwen3 on Hugging Face Spaces.")

    chatbot = gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            gr.Textbox(
                label="System prompt",
                value="You are a helpful assistant.",
                lines=2,
            ),
            gr.Slider(64, 2048, value=MAX_NEW_TOK, step=64, label="Max new tokens"),
            gr.Slider(0.0, 2.0,  value=0.7,        step=0.05, label="Temperature"),
            gr.Slider(0.1, 1.0,  value=0.9,        step=0.05, label="Top-p"),
        ],
        additional_inputs_accordion=gr.Accordion("⚙️ Settings", open=False),
        examples=[
            ["Explain transformers in simple terms."],
            ["Write a haiku about machine learning."],
            ["What is the capital of France?"],
        ],
        cache_examples=False,
    )

demo.queue(max_size=5)
demo.launch()