import gradio as gr
import torch
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import os
import time

MODEL_PATH  = os.getenv("MODEL_PATH", "./model")
MODEL_NAME  = "distilbert-base-multilingual-cased"
MAX_LENGTH  = 512
DEVICE      = torch.device("cuda" if torch.cuda.is_available() else "cpu")
TEMPERATURE = 2.0

print(f"Loading model dari: {MODEL_PATH}")
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model     = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
    model.to(DEVICE)
    model.eval()
    print(f"Model berhasil dimuat! Device: {DEVICE} | Temperature: {TEMPERATURE}")
except Exception as e:
    print(f"Model lokal tidak ditemukan, fallback ke HuggingFace: {e}")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model     = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)
    model.to(DEVICE)
    model.eval()


def detect_text(text: str):
    text = text.strip()

    if not text:
        return None, None, "⚠️ Teks tidak boleh kosong.", ""
    if len(text) < 20:
        return None, None, "⚠️ Teks terlalu pendek, minimal 20 karakter.", ""
    if len(text) > 10000:
        return None, None, "⚠️ Teks terlalu panjang, maksimal 10.000 karakter.", ""

    start_time = time.time()

    inputs = tokenizer(
        text,
        max_length     = MAX_LENGTH,
        padding        = "max_length",
        truncation     = True,
        return_tensors = "pt"
    )
    inputs = {k: v.to(DEVICE) for k, v in inputs.items()
              if k in ["input_ids", "attention_mask"]}

    with torch.no_grad():
        logits        = model(**inputs).logits
        scaled_logits = logits / TEMPERATURE
        probs         = torch.softmax(scaled_logits, dim=-1).cpu().numpy()[0]

    prob_human   = float(probs[0])
    prob_ai      = float(probs[1])
    predicted    = "AI" if prob_ai > prob_human else "Human"
    confidence   = max(prob_ai, prob_human)
    process_time = round(time.time() - start_time, 3)

    word_count = len(text.split())
    char_count = len(text)

    label_html = f"""
    <div style="
        background: {'#1a0a0e' if predicted == 'AI' else '#0a1a12'};
        border: 1.5px solid {'#ff4d6d' if predicted == 'AI' else '#00e5b0'};
        border-radius: 14px;
        padding: 24px 28px;
        margin-bottom: 8px;
    ">
        <div style="display:flex; align-items:center; justify-content:space-between; flex-wrap:wrap; gap:12px;">
            <div style="display:flex; align-items:center; gap:14px;">
                <span style="font-size:32px;">{'🤖' if predicted == 'AI' else '👤'}</span>
                <div>
                    <div style="
                        font-family: 'Segoe UI', sans-serif;
                        font-size: 28px;
                        font-weight: 800;
                        color: {'#ff4d6d' if predicted == 'AI' else '#00e5b0'};
                        letter-spacing: -1px;
                        line-height: 1;
                    ">{'Teks AI' if predicted == 'AI' else 'Teks Manusia'}</div>
                    <div style="color:#888; font-size:13px; margin-top:4px;">
                        {'Kemungkinan besar ditulis oleh AI' if predicted == 'AI' else 'Kemungkinan besar ditulis oleh manusia'}
                    </div>
                </div>
            </div>
            <div style="
                background: {'rgba(255,77,109,0.15)' if predicted == 'AI' else 'rgba(0,229,176,0.15)'};
                border: 1px solid {'#ff4d6d' if predicted == 'AI' else '#00e5b0'};
                color: {'#ff4d6d' if predicted == 'AI' else '#00e5b0'};
                padding: 8px 16px;
                border-radius: 8px;
                font-family: monospace;
                font-size: 14px;
                font-weight: 600;
            ">{confidence*100:.1f}% yakin</div>
        </div>

        <div style="margin-top:20px;">
            <div style="display:flex; justify-content:space-between; margin-bottom:6px;">
                <span style="color:#888; font-size:12px; font-family:monospace;">Kemungkinan AI</span>
                <span style="color:#ff4d6d; font-size:12px; font-family:monospace; font-weight:600;">{prob_ai*100:.1f}%</span>
            </div>
            <div style="background:#1e1e2e; border-radius:99px; height:8px; overflow:hidden;">
                <div style="
                    background: linear-gradient(90deg, #ff4d6d, #ff8fa3);
                    width: {prob_ai*100}%;
                    height: 100%;
                    border-radius: 99px;
                "></div>
            </div>
        </div>

        <div style="margin-top:14px;">
            <div style="display:flex; justify-content:space-between; margin-bottom:6px;">
                <span style="color:#888; font-size:12px; font-family:monospace;">Kemungkinan Manusia</span>
                <span style="color:#00e5b0; font-size:12px; font-family:monospace; font-weight:600;">{prob_human*100:.1f}%</span>
            </div>
            <div style="background:#1e1e2e; border-radius:99px; height:8px; overflow:hidden;">
                <div style="
                    background: linear-gradient(90deg, #00e5b0, #00ffd5);
                    width: {prob_human*100}%;
                    height: 100%;
                    border-radius: 99px;
                "></div>
            </div>
        </div>

        <div style="
            display:grid; grid-template-columns: repeat(3,1fr);
            border-top: 1px solid #2a2a3e;
            margin-top: 20px;
            padding-top: 16px;
            gap: 0;
        ">
            <div style="text-align:center; border-right:1px solid #2a2a3e; padding: 8px;">
                <div style="font-family:monospace; font-size:20px; font-weight:600; color:#e8e8f0;">{word_count:,}</div>
                <div style="font-size:11px; color:#666; text-transform:uppercase; letter-spacing:0.5px; margin-top:2px;">Kata</div>
            </div>
            <div style="text-align:center; border-right:1px solid #2a2a3e; padding: 8px;">
                <div style="font-family:monospace; font-size:20px; font-weight:600; color:#e8e8f0;">{char_count:,}</div>
                <div style="font-size:11px; color:#666; text-transform:uppercase; letter-spacing:0.5px; margin-top:2px;">Karakter</div>
            </div>
            <div style="text-align:center; padding: 8px;">
                <div style="font-family:monospace; font-size:20px; font-weight:600; color:#e8e8f0;">{process_time}s</div>
                <div style="font-size:11px; color:#666; text-transform:uppercase; letter-spacing:0.5px; margin-top:2px;">Waktu Proses</div>
            </div>
        </div>
    </div>
    """

    return label_html, None, "", ""


# ── Custom CSS
css = """
@import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;700;800&family=DM+Mono:wght@400;500&family=DM+Sans:wght@300;400;500&display=swap');

body, .gradio-container {
    background: #0a0a0f !important;
    font-family: 'DM Sans', sans-serif !important;
}

.gradio-container {
    max-width: 860px !important;
    margin: 0 auto !important;
}

/* Header */
#header {
    text-align: center;
    padding: 48px 0 32px;
    border-bottom: 1px solid rgba(255,255,255,0.06);
    margin-bottom: 32px;
}

#header h1 {
    font-family: 'Syne', sans-serif !important;
    font-size: clamp(32px, 5vw, 52px) !important;
    font-weight: 800 !important;
    letter-spacing: -2px !important;
    color: #e8e8f0 !important;
    line-height: 1.05 !important;
    margin: 0 !important;
}

#header p {
    color: #6b6b80 !important;
    font-size: 15px !important;
    margin-top: 12px !important;
    font-weight: 300 !important;
}

/* Pills */
#pills {
    display: flex;
    justify-content: center;
    gap: 8px;
    flex-wrap: wrap;
    margin: 20px 0 0;
}

/* Textarea */
textarea {
    background: #12121a !important;
    border: 1px solid rgba(255,255,255,0.07) !important;
    border-radius: 14px !important;
    color: #e8e8f0 !important;
    font-family: 'DM Sans', sans-serif !important;
    font-size: 15px !important;
    line-height: 1.7 !important;
    padding: 20px !important;
    resize: none !important;
}

textarea:focus {
    border-color: rgba(255,255,255,0.15) !important;
    outline: none !important;
    box-shadow: none !important;
}

textarea::placeholder { color: #4a4a5a !important; }

/* Button */
#scan-btn {
    background: #00e5b0 !important;
    color: #0a0a0f !important;
    border: none !important;
    border-radius: 10px !important;
    font-family: 'Syne', sans-serif !important;
    font-weight: 700 !important;
    font-size: 14px !important;
    letter-spacing: 0.3px !important;
    padding: 12px 28px !important;
    cursor: pointer !important;
    transition: all 0.2s !important;
    width: 100% !important;
}

#scan-btn:hover {
    background: #00ffbf !important;
    transform: translateY(-1px) !important;
}

/* Hide default Gradio labels */
.gr-label { display: none !important; }
label > span { display: none !important; }

/* Error message */
#error-box textarea {
    color: #ff4d6d !important;
    background: rgba(255,77,109,0.06) !important;
    border-color: rgba(255,77,109,0.2) !important;
    font-family: 'DM Mono', monospace !important;
    font-size: 13px !important;
}

/* Footer */
#footer {
    text-align: center;
    padding: 32px 0;
    border-top: 1px solid rgba(255,255,255,0.06);
    margin-top: 40px;
    color: #3a3a4a;
    font-family: 'DM Mono', monospace;
    font-size: 11px;
}
"""

# ── Build UI
with gr.Blocks(css=css, title="TextScan AI — Deteksi Teks AI") as demo:

    gr.HTML("""
    <div id="header">
        <h1>Apakah teks ini ditulis oleh <span style="background:linear-gradient(90deg,#00e5b0,#0090ff);-webkit-background-clip:text;-webkit-text-fill-color:transparent;background-clip:text;">AI</span> atau manusia?</h1>
        <p>Analisis teks menggunakan model DistilBERT yang dilatih untuk membedakan tulisan manusia dan AI</p>
        <div id="pills" style="display:flex;justify-content:center;gap:8px;flex-wrap:wrap;margin-top:20px;">
            <span style="font-family:monospace;font-size:11px;padding:4px 12px;border-radius:20px;border:1px solid #00e5b0;color:#00e5b0;background:rgba(0,229,176,0.08);">DistilBERT ✓ terbaik</span>
            <span style="font-family:monospace;font-size:11px;padding:4px 12px;border-radius:20px;border:1px solid rgba(255,255,255,0.1);color:#6b6b80;">mBERT</span>
            <span style="font-family:monospace;font-size:11px;padding:4px 12px;border-radius:20px;border:1px solid rgba(255,255,255,0.1);color:#6b6b80;">IndoBERT</span>
            <span style="font-family:monospace;font-size:11px;padding:4px 12px;border-radius:20px;border:1px solid rgba(255,255,255,0.1);color:#6b6b80;">ID + EN</span>
        </div>
    </div>
    """)

    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(
                placeholder="Paste atau ketik teks di sini untuk dianalisis...\n\nMinimal 20 karakter, maksimal 10.000 karakter.",
                lines=8,
                max_lines=20,
                label="",
                show_label=False,
            )
            scan_btn = gr.Button("Analisis Teks →", elem_id="scan-btn", variant="primary")

    error_box = gr.Textbox(visible=False, elem_id="error-box", show_label=False)
    result_html = gr.HTML(label="")
    dummy = gr.HTML(visible=False)

    def run_detection(text):
        label_html, _, error, _ = detect_text(text)
        if error:
            return gr.update(value=error, visible=True), ""
        return gr.update(visible=False), label_html

    scan_btn.click(
        fn=run_detection,
        inputs=text_input,
        outputs=[error_box, result_html]
    )

    text_input.submit(
        fn=run_detection,
        inputs=text_input,
        outputs=[error_box, result_html]
    )

    gr.HTML("""
    <div id="footer">
        © 2025 TextScanAI · Skripsi Analisis Komparasi Deteksi Teks AI · BERT
        <span style="margin:0 8px;">·</span>
        DistilBERT · mBERT · IndoBERT
    </div>
    """)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)