import gradio as gr
import torch
import spaces
import re
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread

model_name = "Pyefuri/Qwen2.5-3B-Bahasa-Biak-Final"

print("Memuat tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_name)

print("Memuat model Bahasa Biak...")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    dtype=torch.float16,
    device_map="auto"
)

# =====================================================================
# MODIFIKASI 1: Fungsi interseptor sapaan yang lebih fleksibel (Regex)
# =====================================================================
def cek_sapaan_otomatis(message):
    # Menghapus tanda baca dan mengubah ke huruf kecil
    teks_bersih = re.sub(r'[^\w\s]', '', message.strip().lower())
    
    # Deteksi variasi sapaan dasar secara parsial (termasuk salah ketik seperti hala/helo)
    apakah_menyapa = any(kata in teks_bersih for kata in ["halo", "hala", "hai", "helo", "hi", "p "])
    
    # Deteksi waktu
    apakah_malam = "malam" in teks_bersih
    apakah_pagi = "pagi" in teks_bersih
    apakah_siang = "siang" in teks_bersih
    apakah_sore = "sore" in teks_bersih
    
    # Logika penentu respon instan Bahasa Indonesia
    if apakah_malam or (apakah_menyapa and apakah_malam):
        return "Halo, selamat malam juga! Ada yang bisa saya bantu untuk diterjemahkan ke Bahasa Biak?"
    elif apakah_pagi or (apakah_menyapa and apakah_pagi):
        return "Halo, selamat pagi juga! Ada yang bisa saya bantu hari ini?"
    elif apakah_siang or (apakah_menyapa and apakah_siang):
        return "Halo, selamat siang juga! Ada yang bisa saya bantu?"
    elif apakah_sore or (apakah_menyapa and apakah_sore):
        return "Halo, selamat sore juga! Ada yang bisa saya bantu?"
    elif apakah_menyapa or teks_bersih == "p":
        return "Halo! Ada yang bisa saya bantu untuk diterjemahkan ke Bahasa Biak?"
        
    return None

# Fungsi utama penerjemah dan chat di GPU
@spaces.GPU
def prediksi_terjemahan(message, history):
    respon_sapaan = cek_sapaan_otomatis(message)
    if respon_sapaan:
        yield respon_sapaan
        return

    percakapan = [
        {
            "role": "system", 
            "content": (
                "Kamu adalah asisten AI yang ramah. Kamu menguasai Bahasa Indonesia dan Bahasa Biak. "
                "Aturan interaksi:\n"
                "1. Jika pengguna menyapa (seperti 'halo', 'selamat malam') atau bertanya umum, balaslah dengan ramah menggunakan Bahasa Indonesia.\n"
                "2. Jika pengguna meminta terjemahan atau memasukkan kalimat untuk diartikan ke Bahasa Biak, berikan hasil terjemahan yang akurat."
            )
        }
    ]
    
    # Memproses history format Gradio 6.0 (List of Dictionaries)
    if history:
        for interaksi in history:
            if isinstance(interaksi, dict):
                role = interaksi.get("role")
                content = interaksi.get("content")
                
                # --- BERIKUT ADALAH PENAMBAHAN MODUL PEMBERSIH UNTUK GRADIO 6.0 ---
                if isinstance(content, list):
                    teks_murni = ""
                    for item in content:
                        if isinstance(item, dict) and "text" in item:
                            teks_murni += item["text"]
                    content = teks_murni
                # ------------------------------------------------------------------
                
                if role and content:
                    percakapan.append({"role": role, "content": str(content)})
            elif isinstance(interaksi, (list, tuple)) and len(interaksi) == 2:
                human, assistant = interaksi
                if human:
                    percakapan.append({"role": "user", "content": str(human)})
                if assistant:
                    percakapan.append({"role": "assistant", "content": str(assistant)})
    
    percakapan.append({"role": "user", "content": message})
    
    input_text = tokenizer.apply_chat_template(
        percakapan,
        tokenize=False,
        add_generation_prompt=True
    )
    
    inputs = tokenizer([input_text], return_tensors="pt").to(model.device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    # =====================================================================
    # MODIFIKASI 2: Penyesuaian hyperparameter agar Qwen lebih luwes chitchat
    # =====================================================================
    generation_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=256,
        temperature=0.6,      # Diubah ke 0.6 agar tidak terlalu kaku saat mengobrol biasa
        top_p=0.9,            # Diubah ke 0.9 untuk memberikan variasi kosakata dasar yang lebih natural
        repetition_penalty=1.1, # Diturunkan ke 1.1 agar kalimat mengalir lebih baik
        do_sample=True,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    text_terjemahan = ""
    for new_text in streamer:
        text_terjemahan += new_text
        yield text_terjemahan

# ========================================================
#  UI/UX CONFIGURATION (COMPATIBLE WITH GRADIO 6.0)
# ========================================================
desain_kustom_css = """
#wadah-lebar-chat {
    max-width: 800px !important;
    margin: 0 auto !important;
    float: none !important;
    padding-top: 20px;
}
.judul-pusat {
    text-align: center !important;
    margin-bottom: 15px !important;
}
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="wadah-lebar-chat"):
        gr.Markdown("# 🤖 AI Bahasa Indonesia-Biak", elem_classes="judul-pusat")
        gr.ChatInterface(
            fn=prediksi_terjemahan
        )

if __name__ == "__main__":
    demo.launch(
        theme=gr.themes.Soft(),
        css=desain_kustom_css
    )