import gradio as gr
import torch
import json
import random
import pickle
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# --- 1. CONFIGURATION ---
MODEL_PATH = "model-bert/"
LE_PATH = "model-bert/label_encoder.pkl"
DATASET_PATH = "dataset_itebis_final2.json"
LOGO_URL = "https://itebisdewantara.ac.id/wp-content/uploads/2025/07/Logo-ITEBIS-PGRI-Dewantara-01-1024x536.png"

# --- 2. LOAD MODEL & DATA ---
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH).to(device)
model.eval()

with open(LE_PATH, 'rb') as f:
    le = pickle.load(f)

with open(DATASET_PATH, 'r') as f:
    intents = json.load(f)

# --- 3. LOGIC ---
def chatbot_response(message, history):
    inputs = tokenizer(message, return_tensors="pt", truncation=True, max_length=128).to(device)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
        prob, pred = torch.max(probs, dim=-1)

    tag = le.inverse_transform([pred.item()])[0]
    if prob.item() < 0.3:
        return "Maaf, saya belum memahami pertanyaan tersebut. Bisa coba tanyakan dengan kalimat lain?"

    for intent in intents['intents']:
        if intent['tag'] == tag:
            return random.choice(intent['responses'])
    return "Maaf, jawaban tidak ditemukan."

# --- 4. CUSTOM CSS (PINK USER BUBBLES & COMPACT DESIGN) ---
custom_css = """
/* Menghilangkan padding default Gradio agar bisa Full Screen */
footer {display: none !important;}
.gradio-container {
    max-width: 100% !important;
    margin: 0 !important;
    padding: 0 !important;
    height: 100vh !important;
    display: flex;
    flex-direction: column;
}

/* Header yang tetap di atas */
.header-wrapper {
    display: flex;
    align-items: center;
    justify-content: center;
    gap: 15px;
    padding: 10px;
    background: white;
    border-bottom: 2px solid #800000;
    box-shadow: 0 2px 5px rgba(0,0,0,0.1);
    z-index: 10;
}
.header-wrapper img { width: 80px; }

/* Mengatur container chat agar mengisi sisa layar */
#chat-container {
    flex-grow: 1 !important;
    display: flex;
    flex-direction: column;
    overflow: hidden;
    /* Background ala WhatsApp */
    background-color: #e5ddd5 !important;
    background-image: url('https://user-images.githubusercontent.com/15075759/28719144-86dc0f70-73b1-11e7-911d-60d70fcded21.png') !important;
    background-repeat: repeat !important;
}

/* Memperbaiki tinggi ChatInterface */
.gradio-container .main {
    flex-grow: 1;
    display: flex;
    flex-direction: column;
}

/* Balon Bot */
.message.bot {
    background-color: #ffffff !important;
    color: #000000 !important;
    border-radius: 12px 12px 12px 0px !important;
    border: none !important;
    box-shadow: 0 1px 2px rgba(0,0,0,0.1) !important;
}

/* Balon User (Warna Hijau WhatsApp / Pink sesuai request awal) */
.message.user {
    background-color: #fce7f3 !important; /* pink WA, ganti ke #fce7f3 jika ingin tetap hijau */
    color: #000000 !important;
    border-radius: 12px 12px 0px 12px !important;
    border: none !important;
    box-shadow: 0 1px 2px rgba(0,0,0,0.1) !important;
}

/* Container Input di bagian bawah */
#input-box {
    background: #f0f0f0 !important;
    padding: 10px !important;
    border-top: 1px solid #ccc !important;
}

/* Menyembunyikan elemen Markdown footer agar bersih */
.footer-text {
    padding: 5px;
    background: #f0f0f0;
    text-align: center;
}
"""

# --- 5. INTERFACE ---
with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="red"), title="ITEBIS Assistant") as demo:

    with gr.Group():
        gr.HTML(f"""
            <div class="header-wrapper">
                <img src="{LOGO_URL}" alt="Logo ITEBIS">
                <div class="header-text-container" style="border-left: 3px solid #800000; padding-left: 15px;">
                    <h1 style="color: #800000; font-weight: 800; margin: 0; font-size: 1.2rem;">
                        ITEBIS <span style="color: #f59e0b;">SMART ASSISTANT</span>
                    </h1>
                    <p style="color: #4b5563; margin: 0; font-size: 0.7rem;">Layanan Informasi Akademik Digital</p>
                </div>
            </div>
        """)

    with gr.Column(elem_id="chat-container"):
        gr.ChatInterface(
            fn=chatbot_response,
            examples=["Komplain Nilai","Terlambat Bayar Kuliah","Alur Dispensasi","KTM Hilang","Daftar Bimbingan Skripsi","Syarat Bimbingan Skripsi","Pendaftaran Ujian Hasil", "Pendaftaran Proposal","Lapor Kelulusan","Syarat Yudisium", "Syarat Wisuda","Surat Keterangan Lulus","Ijazah"],
            cache_examples=False,
            # fill_height=True sangat penting untuk tampilan full screen di Gradio 4+
            fill_height=True,
            textbox=gr.Textbox(
                placeholder="Ketik pertanyaan di sini...",
                container=False,
                elem_id="input-box"
            )
        )

    gr.HTML("<div class='footer-text' style='color: #800000; font-size: 0.7rem;'>© 2026 ITEBIS PGRI DEWANTARA</div>")

if __name__ == "__main__":
    demo.launch(share=True)