import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import os
from huggingface_hub import login

# --- Konfigurasi ---
# Gunakan ID model yang spesifik dari Hugging Face Hub
model_id = "google/gemma-3n-E4B-it-litert-preview"
# Dapatkan token dari Hugging Face Space Secrets
# Anda harus mengatur ini di pengaturan Space Anda!
hf_token = os.getenv("HF_TOKEN")

# Lakukan login menggunakan token jika tersedia
if hf_token:
    print("Melakukan login ke Hugging Face Hub...")
    login(token=hf_token)
else:
    print("Peringatan: Secret HF_TOKEN tidak ditemukan. Proses unduh model mungkin gagal jika repo bersifat privat/gated.")

# --- Muat Model dan Tokenizer ---
# Proses ini hanya berjalan sekali saat aplikasi dimulai
print("Memuat tokenizer...")
# Menambahkan trust_remote_code=True untuk mengizinkan kode kustom dari model.
# Token tidak perlu dilewatkan lagi setelah login().
tokenizer = AutoTokenizer.from_pretrained(
    model_id, 
    trust_remote_code=True
)

print("Memuat model... Ini mungkin memakan waktu beberapa menit.")
# Menambahkan trust_remote_code=True juga untuk model.
# torch_dtype=torch.bfloat16 mengurangi penggunaan memori.
# device_map="auto" akan secara otomatis menggunakan GPU jika tersedia.
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
)
print("Model berhasil dimuat!")

# --- Fungsi Inti untuk Chat ---
def chat_function(message, history):
    """
    Fungsi ini dipanggil setiap kali pengguna mengirim pesan.
    Ia memformat input, menghasilkan respons dari model, dan mengembalikan hasilnya.
    """
    # Template chat untuk Gemma memerlukan format spesifik dengan token khusus.
    # Kita akan membangun prompt dengan format ini dari riwayat percakapan.
    chat_template = []
    for user_msg, assistant_msg in history:
        chat_template.append({"role": "user", "content": user_msg})
        chat_template.append({"role": "assistant", "content": assistant_msg})
    
    # Tambahkan pesan pengguna yang baru
    chat_template.append({"role": "user", "content": message})
    
    # Terapkan template chat ke tokenizer
    # Ini akan membuat string prompt yang siap digunakan oleh model
    prompt = tokenizer.apply_chat_template(chat_template, tokenize=False, add_generation_prompt=True)

    # Tokenisasi prompt yang sudah diformat
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    print("\n--- Menghasilkan Respons ---")
    print(f"Prompt yang dikirim ke model:\n{prompt}")

    # Hasilkan respons dari model
    outputs = model.generate(
        **inputs,
        max_new_tokens=1500,  # Batas maksimum token untuk respons
        eos_token_id=tokenizer.eos_token_id  # Berhenti saat token end-of-sentence tercapai
    )
    
    # Decode token output menjadi teks
    # Kita hanya mengambil bagian respons yang baru dihasilkan
    response_text = tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True)
    
    print(f"Respons dari model:\n{response_text}")
    return response_text

# --- Membuat Antarmuka Gradio ---
# gr.ChatInterface adalah komponen UI tingkat tinggi yang menyediakan fungsionalitas chat
chatbot_ui = gr.ChatInterface(
    fn=chat_function,
    title="🤖 Gemma 3n Chat",
    description="""
    Ini adalah antarmuka chat untuk model **google/gemma-3n-E4B-it-litert-preview**. 
    Model ini adalah model besar, jadi harap bersabar untuk responsnya jika berjalan di CPU.
    **Sangat disarankan untuk menggunakan hardware GPU untuk performa terbaik.**
    """,
    examples=[
        ["Ceritakan tentang sejarah singkat Indonesia"],
        ["Buatkan saya resep untuk membuat nasi goreng spesial"],
        ["Jelaskan konsep machine learning dalam 3 paragraf"]
    ],
    cache_examples=False # Nonaktifkan cache untuk contoh dinamis
).queue() # .queue() membuat aplikasi lebih responsif

# --- Menjalankan Aplikasi ---
if __name__ == "__main__":
    chatbot_ui.launch()