import os
import torch
import gradio as gr
from transformers import pipeline

access_token = os.environ.get('hf_token')

if torch.cuda.is_available():
    device_arg = 0  # nvidia gpu
elif torch.backends.mps.is_available():
    device_arg = "mps"  # apple silicon gpu
else:
    device_arg = -1  # cpu fallback

# Initialize text classification pipeline
classifier = pipeline(
    task="text-classification",
    model="Falconsai/offensive_speech_detection",
    device=device_arg,
    return_all_scores=True  # Asks pipeline for both labels
)

# Human-readable mapping for model output
LABEL_MAPPING = {
    "LABEL_0": "Normal Speech",
    "LABEL_1": "Offensive Speech"
}

def detect_offensive_speech(text):
    if not text.strip():
        return {"Empty Input": 1.0}
    
    raw_output = classifier(text)
    
    # --- DEFENSIVE PARSING BLOCK ---
    # Case 1: Nested list [[{...}, {...}]] -> Extract inner list
    if isinstance(raw_output, list) and len(raw_output) > 0 and isinstance(raw_output[0], list):
        items = raw_output[0]
    # Case 2: Flat list [{...}, {...}] -> Use as-is
    elif isinstance(raw_output, list):
        items = raw_output
    # Case 3: Single dictionary {...} -> Wrap in a list
    elif isinstance(raw_output, dict):
        items = [raw_output]
    else:
        return {"Parsing Error": 1.0}
    # -------------------------------

    formatted_results = {}
    for item in items:
        # Extra safety: Ensure 'item' is a dictionary before reading keys
        if isinstance(item, dict) and "label" in item and "score" in item:
            clean_label = LABEL_MAPPING.get(item["label"], item["label"])
            formatted_results[clean_label] = float(item["score"])
            
    # Fallback if parsing completely fails to find labels
    if not formatted_results:
        return {"Unknown Error": 1.0}
        
    return formatted_results

# Define Gradio Interface
demo = gr.Interface(
    fn=detect_offensive_speech,
    inputs=gr.Textbox(placeholder="Enter text here...", lines=3),
    outputs=gr.Label(label="Classification Results", num_top_classes=2),
    title="Offensive Speech Detection",
    description="A local, optimized text-classification pipeline running on Hugging Face weights."
)

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