import re
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


URL_PATTERN = re.compile(r"https?://\S+|www\.\S+", flags=re.IGNORECASE)

def fast_detect_unknown(text: str) -> bool:
    s = text.strip()
    if not s:
        # Empty input is treated as non-language.
        return True
    if URL_PATTERN.search(s):
        # URLs should map to UNKNOWN.
        return True

    total = len(s)
    alpha = sum(ch.isalpha() for ch in s)
    digits = sum(ch.isdigit() for ch in s)
    spaces = sum(ch.isspace() for ch in s)
    symbols = total - alpha - digits - spaces
    non_space = max(1, total - spaces)

    # Mostly numeric strings (ids, phone numbers, etc.).
    if digits / non_space >= 0.8:
        return True
    # Symbol-heavy text is usually not valid language content.
    if symbols / non_space >= 0.45:
        return True
    # Very low alphabetic ratio indicates gibberish-like input.
    if total >= 6 and (alpha / non_space) < 0.2:
        return True
    # Long compact mixed tokens often represent hashes/usernames/keys.
    if " " not in s and total >= 12 and (alpha / non_space) < 0.45 and (digits > 0 or symbols > 0):
        return True
    return False

def predict_with_automodel(
    text: str,
    model: AutoModelForSequenceClassification,
    tokenizer: AutoTokenizer,
    max_length: int = 128,
) -> tuple[str, float]:
    encoded = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=max_length,
    ).to(model.device)
    with torch.no_grad():
        logits = model(**encoded).logits
        probs = torch.softmax(logits, dim=-1).squeeze(0)
    pred_id = int(torch.argmax(probs).item())
    pred_score = float(probs[pred_id].item())
    pred_label = model.config.id2label[pred_id]
    return pred_label, pred_score


if __name__ == "__main__":
    model_dir = "chiennv/langid-mmbert-small"
    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForSequenceClassification.from_pretrained(model_dir, dtype=torch.float16, device_map="auto")
    model.eval()
    texts = [
        "This is a test sentence.",
        "Bonjour tout le monde",
        "https://example123.com/p/9999",
        "Christiano Ronaldo",
        "1234567890",
        "1234567890 is a phone number",
        "Hello",
        "A footballer",
        "a dog is barking",
        "le chatte",
        "le chatte est un animal",
        "고양이는 동물입니다.",
        "猫は動物です。",
    ]
    for t in texts:
        if fast_detect_unknown(t):
            print(t, "->", "UNKNOWN", 1.0)
            continue
        label, score = predict_with_automodel(t, model=model, tokenizer=tokenizer)
        print(t, "->", label, round(score, 4))