import gradio as gr
from transformers import T5Tokenizer, T5ForConditionalGeneration # Import T5Tokenizer dan T5ForConditionalGeneration
import torch

# ---------- LOAD MODEL DAN TOKENIZER ----------
try:
    # Memuat tokenizer dan model secara terpisah untuk model T5
    tokenizer = T5Tokenizer.from_pretrained("cahya/t5-base-indonesian-summarization-cased")
    model = T5ForConditionalGeneration.from_pretrained("cahya/t5-base-indonesian-summarization-cased")
    
    # Pindahkan model ke GPU jika tersedia
    device = 0 if torch.cuda.is_available() else -1
    if device != -1:
        model.to(f"cuda:{device}")
    
    print("Model dan tokenizer Bahasa Indonesia berhasil dimuat.")
except Exception as e:
    tokenizer = None
    model = None
    print(f"Error saat memuat model dan tokenizer: {str(e)}") # Cetak error ke konsol jika gagal dimuat

# ---------- SUMMARIZATION FUNCTION ----------
def summarize_text_simple(text_input, min_length_val=30, max_length_val=150):
    # Cek apakah model dan tokenizer berhasil dimuat
    if tokenizer is None or model is None:
        return "❌ Error: Model ringkasan gagal dimuat. Coba lagi nanti."

    if not text_input.strip():
        return "⚠️ Mohon masukkan teks yang ingin diringkas!"

    # Pastikan panjang minimum dan maksimum masuk akal
    if min_length_val >= max_length_val:
        return "⚠️ Panjang minimum harus lebih kecil dari panjang maksimum!"
    if min_length_val <= 0 or max_length_val <= 0:
        return "⚠️ Panjang tidak boleh nol atau negatif!"
    
    try:
        # Menambahkan prefix "summarize: " yang umum digunakan untuk T5 summarization
        # Juga menambahkan truncation=True di tokenizer
        input_ids = tokenizer.encode("summarize: " + text_input, 
                                     return_tensors="pt", 
                                     max_length=512, # Batasi panjang input token T5 (umumnya 512)
                                     truncation=True)
        
        # Pindahkan input_ids ke GPU jika model ada di GPU
        if device != -1:
            input_ids = input_ids.to(f"cuda:{device}")

        # Lakukan generasi ringkasan
        summary_ids = model.generate(
            input_ids,
            min_length=int(min_length_val),
            max_length=int(max_length_val),
            num_beams=4, # Jumlah beam untuk beam search (meningkatkan kualitas)
            early_stopping=True # Hentikan generasi lebih awal jika semua beam selesai
        )
        
        # Dekode token hasil menjadi teks
        summarized_text = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
        
        result_message = f"""
        <h3>Teks Ringkasan Anda:</h3>
        <p>{summarized_text}</p>
        """
        
        return result_message
    
    except Exception as e:
        return f"❌ Terjadi kesalahan saat meringkas: {str(e)}"

# ---------- GRADIO INTERFACE ----------
with gr.Blocks(title="Aplikasi Ringkasan Teks Sederhana (ID)") as demo:
    gr.Markdown("# 📝 Aplikasi Ringkasan Teks Sederhana (Bahasa Indonesia)")
    gr.Markdown("Masukkan teks panjang berbahasa Indonesia di bawah ini untuk mendapatkan versi ringkasnya.")

    with gr.Row():
        text_input = gr.Textbox(
            label="Teks Asli (Bahasa Indonesia)",
            placeholder="Masukkan teks panjang berbahasa Indonesia yang ingin Anda ringkas di sini...",
            lines=10
        )
    
    with gr.Row():
        min_length_slider = gr.Slider(
            minimum=10,
            maximum=100,
            value=30,
            step=1,
            label="Panjang Ringkasan Minimum"
        )
        max_length_slider = gr.Slider(
            minimum=50,
            maximum=200, # Batasi maksimum yang lebih masuk akal untuk ringkasan
            value=80,
            step=1,
            label="Panjang Ringkasan Maksimum"
        )

    summarize_btn = gr.Button("✨ Ringkas Sekarang")

    summary_output = gr.HTML(label="Hasil Ringkasan")

    # Menghubungkan tombol ke fungsi ringkasan
    summarize_btn.click(
        fn=summarize_text_simple,
        inputs=[text_input, min_length_slider, max_length_slider],
        outputs=summary_output
    )

    gr.Markdown("""
    ---
    <div style='text-align: center; margin-top: 20px;'>
        <p>Didukung oleh Hugging Face Transformers (Model: cahya/t5-base-indonesian-summarization-cased) dan Gradio.</p>
    </div>
    """)

# Jalankan antarmuka
if __name__ == "__main__":
    demo.launch()