"""Inference script for the hate speech classifier."""

import os
import re
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from peft import PeftModel

MODEL_PATH = "./trained/Best_hate_speech"
MAX_LENGTH = 128
LABEL_NAMES = {
    0: "Hate Speech",
    1: "Offensive Language",
    2: "Neither",
}


def normalize_text(text):
    """Light tweet-style normalization (URLs, mentions, hashtags, RT)."""
    if not isinstance(text, str):
        return ""
    text = re.sub(r"https?://\S+", "", text)
    text = re.sub(r"@\w+", "", text)
    text = re.sub(r"#(\w+)", r"\1", text)
    text = re.sub(r"^RT\s+", "", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text


def load_model(model_path=MODEL_PATH, device=None):
    """Load trained model and tokenizer. Supports both LoRA adapters and full models."""
    if device is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"

    tokenizer = AutoTokenizer.from_pretrained(model_path)

    if os.path.exists(os.path.join(model_path, "adapter_config.json")):
        from peft import PeftConfig
        peft_config = PeftConfig.from_pretrained(model_path)
        base_model = AutoModelForSequenceClassification.from_pretrained(
            peft_config.base_model_name_or_path,
            num_labels=len(LABEL_NAMES),
        )
        model = PeftModel.from_pretrained(base_model, model_path)
        model = model.merge_and_unload()
        print(f"LoRA adapter loaded and merged from {model_path}")
    else:
        model = AutoModelForSequenceClassification.from_pretrained(model_path)
        print(f"Full model loaded from {model_path}")

    model.to(device)
    model.eval()
    print(f"Device: {device}")
    return model, tokenizer, device


def predict(text, model, tokenizer, device):
    """Classify a single text. Returns label, confidence, and all probabilities."""
    text = normalize_text(text)
    encoding = tokenizer(
        text,
        max_length=MAX_LENGTH,
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )
    encoding = {k: v.to(device) for k, v in encoding.items()}

    with torch.no_grad():
        logits = model(**encoding).logits
        probs = torch.softmax(logits, dim=-1).squeeze()

    pred_id = int(probs.argmax().item())
    return {
        "label": LABEL_NAMES[pred_id],
        "confidence": probs[pred_id].item(),
        "probabilities": {LABEL_NAMES[i]: probs[i].item() for i in range(len(LABEL_NAMES))},
    }


def predict_batch(texts, model, tokenizer, device, batch_size=32):
    """Classify a list of texts efficiently in batches."""
    results = []
    for i in range(0, len(texts), batch_size):
        batch = [normalize_text(t) for t in texts[i : i + batch_size]]
        encoding = tokenizer(
            batch,
            max_length=MAX_LENGTH,
            padding="max_length",
            truncation=True,
            return_tensors="pt",
        )
        encoding = {k: v.to(device) for k, v in encoding.items()}

        with torch.no_grad():
            logits = model(**encoding).logits
            probs = torch.softmax(logits, dim=-1)

        for j in range(len(batch)):
            p = probs[j]
            pred_id = int(p.argmax().item())
            results.append({
                "text": batch[j],
                "label": LABEL_NAMES[pred_id],
                "confidence": p[pred_id].item(),
                "probabilities": {LABEL_NAMES[k]: p[k].item() for k in range(len(LABEL_NAMES))},
            })
    return results


if __name__ == "__main__":
    model, tokenizer, device = load_model()

    examples = [
        "I hate all of those people, they should be removed from this country.",
        "That movie was terrible, worst film I've seen all year.",
        "You're such an idiot, nobody likes you.",
        "The weather is nice today, I think I'll go for a walk.",
    ]

    print("\n" + "=" * 70)
    print("HATE SPEECH CLASSIFIER")
    print("=" * 70)
    for text in examples:
        result = predict(text, model, tokenizer, device)
        if result["label"] == "Hate Speech":
            icon = "[X]"
        elif result["label"] == "Offensive Language":
            icon = "[!]"
        else:
            icon = "[OK]"
        print(f"\n{icon}  [{result['label']}] ({result['confidence']:.1%})")
        print(f"   \"{text}\"")

    print("\n" + "=" * 70)
    print("Type a sentence to classify (or 'quit' to exit):")
    print("=" * 70)
    while True:
        try:
            user_input = input("\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not user_input or user_input.lower() in ("quit", "exit", "q"):
            break
        result = predict(user_input, model, tokenizer, device)
        if result["label"] == "Hate Speech":
            icon = "[X]"
        elif result["label"] == "Offensive Language":
            icon = "[!]"
        else:
            icon = "[OK]"
        print(f"{icon}  {result['label']} ({result['confidence']:.1%})")
        for name, prob in result["probabilities"].items():
            print(f"   {name}: {prob:.1%}")