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

MODEL_PATH = "." 

try:
    print(f"Loading model from {MODEL_PATH}...")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
    model.eval()
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    print("Falling back to default DistilBERT (untrained) for interface demo purpose only.")
    model_name = "distilbert-base-uncased"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

def predict_toxicity(text):
    if not text or not text.strip():
        return "<h3 style='color: grey;'>Waiting for input...</h3>", {}

    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
    
    with torch.no_grad():
        logits = model(**inputs).logits
        temperature = 2.0 
        probs = torch.softmax(logits / temperature, dim=1)[0].cpu().numpy()

    non_toxic_prob = float(probs[0])
    toxic_prob = float(probs[1])
    toxic_pct = f"{toxic_prob:.2%}"
    normal_pct = f"{non_toxic_prob:.2%}"

    if toxic_prob > 0.5:
        result_html = f"""
        <div style="background-color: #ffcccc; border: 2px solid red; border-radius: 10px; padding: 20px; text-align: center;">
            <h2 style="color: red; margin: 0;">🔴 TOXIC DETECTED</h2>
            <p style="margin: 5px 0 0 0; color: darkred; font-weight: bold;">Confidence: {toxic_pct}</p>
            <p style="margin: 5px 0 0 0; color: #cc0000; font-size: 0.9em;">This comment violates community guidelines.</p>
        </div>
        """
    else:
        result_html = f"""
        <div style="background-color: #ccffcc; border: 2px solid green; border-radius: 10px; padding: 20px; text-align: center;">
            <h2 style="color: green; margin: 0;">🟢 CONTENT IS SAFE</h2>
            <p style="margin: 5px 0 0 0; color: darkgreen; font-weight: bold;">Confidence: {normal_pct}</p>
            <p style="margin: 5px 0 0 0; color: #006600; font-size: 0.9em;">No harmful language detected.</p>
        </div>
        """

    return result_html, {"Toxic": toxic_prob, "Normal": non_toxic_prob}

with gr.Blocks(theme=gr.themes.Soft(), title="AI Toxic Comment Detector") as demo:
    
    gr.Markdown(
        """
        # 🛡️ AI Toxic Comment Detector
        ### Ensuring Digital Safety with Artificial Intelligence
        This tool supports **SDG 16 (Peace, Justice and Strong Institutions)** by reducing online hate speech 
        and **SDG 3 (Good Health and Well-being)** by protecting mental health.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                lines=5, 
                placeholder="Type a comment here to analyze...", 
                label="📝 Input Text",
                info="Supports English text analysis."
            )
            
            with gr.Row():
                clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                submit_btn = gr.Button("🚀 Analyze", variant="primary")

        with gr.Column(scale=1):
            output_html = gr.HTML(label="Verdict")
            output_chart = gr.Label(num_top_classes=2, label="Confidence Score 📊")

    gr.Markdown("### 🧪 Try these examples:")
    examples = gr.Examples(
        examples=[
            ["I really appreciate your perspective on this article, well done!"],
            ["You are absolute garbage, nobody wants you here."],
            ["I love how you just don't care about looking professional."], 
            ["This movie was crazy good! I killed it at the gym today."]
        ],
        inputs=input_text
    )

    submit_btn.click(
        fn=predict_toxicity, 
        inputs=input_text, 
        outputs=[output_html, output_chart]
    )
    
    clear_btn.click(
        lambda: ("", "<h3 style='color: grey;'>Waiting for input...</h3>", {}),
        inputs=None,
        outputs=[input_text, output_html, output_chart]
    )

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