import gradio as gr
import torch
import re
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from deep_translator import GoogleTranslator
from indic_transliteration.sanscript import transliterate, ITRANS, TAMIL, MALAYALAM

# -------------------- DEVICE --------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# -------------------- LOAD ALL MODELS --------------------

models = {}
tokenizers = {}

MODEL_PATHS = {
    "Telugu": "Thilak118/indic-bert-toxicity-classifier",
    "Tamil": "Thilak118/indic-bert-toxicity-classifier_tamil",
    "Kannada": "Thilak118/indic-bert-toxicity-classifier_kannada",
    "Malayalam": "Thilak118/indic-bert-toxicity-classifier_malayalam"
}

for lang, path in MODEL_PATHS.items():
    tokenizers[lang] = AutoTokenizer.from_pretrained(path)
    models[lang] = AutoModelForSequenceClassification.from_pretrained(path).to(device)
    models[lang].eval()

# -------------------- TRANSLATORS --------------------

telugu_translator = GoogleTranslator(source="en", target="te")
kannada_translator = GoogleTranslator(source="en", target="kn")

# -------------------- CLEANING FUNCTIONS --------------------

def clean_text(text, language):
    ranges = {
        "Telugu": r"[^\u0C00-\u0C7F\s.,!?]",
        "Tamil": r"[^\u0B80-\u0BFF\s.,!?]",
        "Kannada": r"[^\u0C80-\u0CFF\s.,!?]",
        "Malayalam": r"[^\u0D00-\u0D7F\s.,!?]"
    }
    text = re.sub(ranges[language], "", str(text))
    text = re.sub(r"\s+", " ", text).strip()
    return text

# -------------------- TRANSLITERATION --------------------

def transliterate_text(text, language):

    if not text or not text.strip():
        return ""

    try:
        if language == "Telugu":
            return telugu_translator.translate(text)

        elif language == "Tamil":
            return transliterate(text, ITRANS, TAMIL)

        elif language == "Kannada":
            return kannada_translator.translate(text)

        elif language == "Malayalam":
            return transliterate(text, ITRANS, MALAYALAM)

    except:
        return "Transliteration failed"

# -------------------- PREDICTION FUNCTION --------------------

def unified_predict(language, input_text):

    script_text = transliterate_text(input_text, language)

    if "failed" in str(script_text).lower():
        return f"{language} Text: {script_text}\nPrediction: Failed"

    cleaned_text = clean_text(script_text, language)

    if not cleaned_text:
        return "Invalid input"

    tokenizer = tokenizers[language]
    model = models[language]

    inputs = tokenizer(
        cleaned_text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=128
    ).to(device)

    with torch.no_grad():
        outputs = model(**inputs)

    probs = torch.softmax(outputs.logits, dim=1)[0]
    prediction = torch.argmax(probs).item()
    confidence = probs[prediction].item() * 100

    label = "Toxic" if prediction == 0 else "Non-Toxic"

    return (
        f"{language} Script Text: {cleaned_text}\n"
        f"Prediction: {label}\n"
        f"Confidence: {confidence:.2f}%"
    )

# -------------------- GRADIO UI --------------------

with gr.Blocks(title="South Indian Languages Toxicity Classifier") as demo:

    gr.Markdown(
        """
        # Unified Toxicity Classifier 🇮🇳  
        Supports:
        - Telugu
        - Tamil
        - Kannada
        - Malayalam
        
        Enter text in English transliteration and select language.
        """
    )

    language_dropdown = gr.Dropdown(
        ["Telugu", "Tamil", "Kannada", "Malayalam"],
        label="Select Language",
        value="Telugu"
    )

    input_text = gr.Textbox(
        label="Enter Text (English Transliteration)",
        placeholder="Type your comment here...",
        lines=2
    )

    predict_btn = gr.Button("Predict Toxicity")

    output_text = gr.Textbox(
        label="Prediction Output",
        interactive=False,
        lines=5
    )

    predict_btn.click(
        fn=unified_predict,
        inputs=[language_dropdown, input_text],
        outputs=output_text
    )

demo.launch()