import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_NAME = "Fatimasane26/distilbert-toxic-jigsaw"
MAX_LENGTH = 128
THRESHOLD  = 0.5

CATEGORIES = [
    ("Toxic",         "☠️", "#f43f5e"),
    ("Severe Toxic",  "💀", "#e11d48"),
    ("Obscene",       "🤬", "#fb923c"),
    ("Threat",        "⚔️", "#a78bfa"),
    ("Insult",        "😤", "#facc15"),
    ("Identity Hate", "🎯", "#f472b6"),
]

print("⏳ Loading model...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model     = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
print(f"✅ Model ready on {device}")


def predict(text):
    if not text or not text.strip():
        return "", "", build_bars(None)

    inputs = tokenizer(
        text, return_tensors="pt", truncation=True,
        max_length=MAX_LENGTH, padding=True,
        return_token_type_ids=False
    ).to(device)

    with torch.no_grad():
        probs = torch.sigmoid(model(**inputs).logits).cpu().numpy()[0]

    detected   = [name for (name, _, __), p in zip(CATEGORIES, probs) if p > THRESHOLD]
    max_score  = float(max(probs))

    if max_score < 0.2:
        verdict    = "✅  Clean"
        confidence = f"{int(max_score*100)}%  — No toxicity detected"
    elif max_score < THRESHOLD:
        verdict    = "⚠️  Low Risk"
        confidence = f"{int(max_score*100)}%  — Potentially sensitive"
    else:
        cats       = ", ".join(detected) if detected else "Unknown"
        verdict    = "🚨  Toxic"
        confidence = f"{int(max_score*100)}%  — {cats}"

    return verdict, confidence, build_bars(probs)


def build_bars(probs):
    if probs is None:
        return """
        <div style='display:flex;align-items:center;justify-content:center;
                    height:220px;color:#475569;font-family:sans-serif;flex-direction:column;gap:10px;'>
            <span style='font-size:36px;'>📊</span>
            <span style='font-size:14px;'>Probabilities will appear here</span>
        </div>"""

    rows = ""
    for (name, emoji, color), score in zip(CATEGORIES, probs):
        pct    = int(score * 100)
        width  = max(pct, 1)
        bold   = "font-weight:700;" if score > THRESHOLD else "opacity:0.6;"
        badge  = f"<span style='background:{color};color:#000;font-size:10px;padding:1px 8px;border-radius:999px;font-weight:800;margin-left:6px;'>!</span>" if score > THRESHOLD else ""
        rows += f"""
        <div style='margin-bottom:13px;'>
          <div style='display:flex;justify-content:space-between;align-items:center;margin-bottom:4px;'>
            <span style='color:#cbd5e1;font-size:13px;{bold}'>{emoji} {name}{badge}</span>
            <span style='color:#94a3b8;font-size:13px;font-weight:700;'>{pct}%</span>
          </div>
          <div style='background:#1e293b;border-radius:999px;height:9px;overflow:hidden;'>
            <div style='width:{width}%;height:100%;background:linear-gradient(90deg,{color}aa,{color});
                        border-radius:999px;transition:width 0.4s;'></div>
          </div>
        </div>"""

    return f"<div style='font-family:system-ui,sans-serif;padding:4px 2px;'>{rows}</div>"


CSS = """
* { box-sizing: border-box; }
body, .gradio-container { background: #0a0f1e !important; font-family: system-ui, sans-serif; }
.gradio-container { max-width: 960px !important; margin: 0 auto !important; padding: 0 16px !important; }
footer { display: none !important; }

/* Card style */
.card {
    background: #0f172a !important;
    border: 1px solid #1e293b !important;
    border-radius: 14px !important;
    padding: 0 !important;
    overflow: hidden !important;
}

/* Label headers like the reference */
.card-label {
    background: linear-gradient(135deg, #6d28d9, #0ea5e9);
    color: white;
    font-weight: 700;
    font-size: 13px;
    padding: 8px 16px;
    border-radius: 999px;
    display: inline-block;
    margin-bottom: 12px;
    letter-spacing: 0.3px;
}

/* Inputs */
textarea, input[type=text] {
    background: #1e293b !important;
    border: 1px solid #334155 !important;
    border-radius: 10px !important;
    color: #e2e8f0 !important;
    font-size: 15px !important;
}
textarea:focus { border-color: #6d28d9 !important; box-shadow: 0 0 0 2px #6d28d920 !important; }

/* Buttons */
button.primary {
    background: linear-gradient(135deg, #6d28d9, #0ea5e9) !important;
    border: none !important; border-radius: 10px !important;
    color: white !important; font-weight: 700 !important;
    font-size: 14px !important; padding: 10px 0 !important;
    transition: opacity 0.2s !important;
}
button.primary:hover { opacity: 0.85 !important; }
button.secondary {
    background: #1e293b !important; border: 1px solid #334155 !important;
    border-radius: 10px !important; color: #94a3b8 !important;
    font-weight: 600 !important;
}

/* Output text boxes */
.output-text input, .output-text textarea {
    background: #1e293b !important; border: 1px solid #334155 !important;
    color: #e2e8f0 !important; font-weight: 600 !important;
    border-radius: 10px !important;
}

/* Examples */
.examples-holder { background: transparent !important; }
table.examples { background: transparent !important; }
table.examples td { color: #64748b !important; font-size: 12px !important; }
table.examples tr:hover td { color: #e2e8f0 !important; background: #1e293b !important; }
"""

EXAMPLES = [
    "You are such an idiot, I hate you so much!",
    "I will find you and destroy everything you love.",
    "Thank you for your help, really appreciate it!",
    "People like you should not be allowed here.",
    "Great discussion, learned a lot today!",
    "I'll make sure you regret this, I promise.",
]

with gr.Blocks(css=CSS, title="🛡️ Toxic Comment Classifier") as demo:

    # ── HEADER ──────────────────────────────────────────────────────
    gr.HTML("""
    <div style='background:linear-gradient(135deg,#1e1b4b 0%,#0c4a6e 100%);
                border-radius:18px;padding:36px 40px;margin-bottom:24px;'>
      <h1 style='color:#f8fafc;font-size:30px;font-weight:800;margin:0 0 8px;'>
        🛡️ Toxic Comment Classifier
      </h1>
      <p style='color:#94a3b8;font-size:15px;margin:0 0 20px;'>
        Analyze any text comment and detect toxicity across 6 categories in real time.
      </p>
      <div style='display:flex;gap:10px;flex-wrap:wrap;'>
        <span style='background:#ffffff18;color:#e2e8f0;padding:5px 14px;border-radius:999px;font-size:12px;font-weight:600;'>🤗 DistilBERT</span>
        <span style='background:#ffffff18;color:#e2e8f0;padding:5px 14px;border-radius:999px;font-size:12px;font-weight:600;'>📦 223K Comments</span>
        <span style='background:#ffffff18;color:#e2e8f0;padding:5px 14px;border-radius:999px;font-size:12px;font-weight:600;'>🏷️ 6 Categories</span>
        <span style='background:#ffffff18;color:#e2e8f0;padding:5px 14px;border-radius:999px;font-size:12px;font-weight:600;'>⚡ Real-time</span>
      </div>
    </div>""")

    # ── MAIN LAYOUT ─────────────────────────────────────────────────
    with gr.Row(equal_height=False):

        # LEFT — Input
        with gr.Column(scale=5):
            gr.HTML("<div class='card-label'>💬 Comment to Analyze</div>")
            text_input = gr.Textbox(
                placeholder="Type or paste a comment here…",
                lines=6, show_label=False, container=False
            )
            with gr.Row():
                btn = gr.Button("🔍  Analyze Now", variant="primary", scale=3)
                gr.ClearButton([text_input], value="✕ Clear", scale=1)

            gr.HTML("<div style='color:#334155;font-size:11px;margin:10px 0 4px;text-align:center;'>— Try an example —</div>")
            gr.Examples(
                examples=[[e] for e in EXAMPLES],
                inputs=text_input, label=""
            )

        # RIGHT — Output
        with gr.Column(scale=5):
            gr.HTML("<div class='card-label'>🎯 Predicted Verdict</div>")
            verdict_out = gr.Textbox(
                show_label=False, interactive=False,
                placeholder="—", container=False,
                elem_classes=["output-text"]
            )

            gr.HTML("<div class='card-label' style='margin-top:16px;'>📊 Confidence</div>")
            conf_out = gr.Textbox(
                show_label=False, interactive=False,
                placeholder="—", container=False,
                elem_classes=["output-text"]
            )

            gr.HTML("<div class='card-label' style='margin-top:16px;'>📈 Category Probabilities</div>")
            bars_out = gr.HTML("""
            <div style='display:flex;align-items:center;justify-content:center;
                        height:180px;color:#334155;font-family:sans-serif;flex-direction:column;gap:8px;'>
                <span style='font-size:32px;'>📊</span>
                <span style='font-size:13px;'>Probabilities will appear here</span>
            </div>""")

    # ── FOOTER ──────────────────────────────────────────────────────
    gr.HTML("""
    <div style='text-align:center;padding:20px 0 4px;color:#1e293b;
                font-size:12px;font-family:sans-serif;'>
        NLP Final Project &nbsp;·&nbsp; Built with HuggingFace Transformers &amp; Gradio
    </div>""")

    btn.click(fn=predict, inputs=text_input, outputs=[verdict_out, conf_out, bars_out])
    text_input.submit(fn=predict, inputs=text_input, outputs=[verdict_out, conf_out, bars_out])

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