import os
import torch
import gradio gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. Load Model dan Tokenizer (Qwen3-0.6B)
MODEL_NAME = "Qwen/Qwen3-0.6B"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype="auto",
    device_map="auto"
)

def respond(message, history, system_message, max_tokens, temperature, top_p):
    # Susun percakapan sesuai format yang disarankan Qwen3
    # Catatan: Best practice Qwen3 menyarankan riwayat chat tidak mengikutsertakan proses berpikir sebelumnya
    messages = [{"role": "system", "content": system_message}]
    
    # Pada Gradio 6, riwayat dikirim dalam format objek khusus atau dictionary.
    # Kode ini mengekstrak konten user dan assistant secara aman.
    for turn in history:
        user_content = turn.get("user") if isinstance(turn, dict) else getattr(turn, "user", None)
        assistant_content = turn.get("assistant") if isinstance(turn, dict) else getattr(turn, "assistant", None)
        
        if user_content:
            messages.append({"role": "user", "content": user_content})
        if assistant_content:
            messages.append({"role": "assistant", "content": assistant_content})
            
    messages.append({"role": "user", "content": message})
    
    # Terapkan template chat bawaan Qwen3 (default: enable_thinking=True)
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=True
    )
    
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    # Generate respons
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        top_k=20,          # Sesuai rekomendasi dokumentasi Qwen3
        do_sample=True     # Wajib True untuk thinking mode agar tidak berulang
    )
    
    # Ambil token baru saja
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
    
    # Parsing konten berfikir berdasarkan token khusus </think> (ID: 151668)
    try:
        index = len(output_ids) - output_ids[::-1].index(151668)
    except ValueError:
        index = 0
        
    thinking_content = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
    final_content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
    
    # Format output agar cantik di UI Gradio
    full_response = ""
    if thinking_content:
        full_response += f"💭 *Thinking Process:*\n```text\n{thinking_content}\n```\n\n"
    full_response += final_content
    
    return full_response

# 2. Desain Antarmuka Gradio (Zela 3.0 UI untuk Gradio 6)
with gr.Blocks() as demo:
    gr.ChatInterface(
        respond,
        title="Zela 3.0 🤖✨",
        description="Asisten pintar bertenaga Qwen3-0.6B dengan fitur *Thinking Mode* bawaan.",
        examples=[
            ["Berapa jumlah huruf 'r' pada kata 'strawberry'?"],
            ["Berikan penjelasan singkat apa itu Black Hole."],
            ["Buat skrip python sederhana untuk menghitung faktorial."]
        ],
        additional_inputs=[
            gr.Textbox(value="Kamu adalah Zela 3.0, asisten AI yang ramah, jenius, dan membantu.", label="System Message"),
            gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max New Tokens"),
            # Nilai default disesuaikan dengan rekomendasi Qwen3 (Thinking Mode = 0.6)
            gr.Slider(minimum=0.1, maximum=1.0, value=0.6, step=0.05, label="Temperature"),
            gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"),
        ]
    )

if __name__ == "__main__":
    # Sesuai standar Gradio 6, tema dipasangkan di fungsi launch()
    demo.launch(theme="soft")