"""Load the original fully fine-tuned SafeAge XLM-R binary baseline."""
import argparse
import contextlib
import json
from pathlib import Path

DEFAULT_REPO = "mehedi4/SafeAge-XLM-R-Base"


def load_classifier(source=DEFAULT_REPO, revision=None, device=None, precision="auto"):
    import torch
    from huggingface_hub import snapshot_download
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    folder = Path(source)
    if not folder.is_dir():
        folder = Path(snapshot_download(
            repo_id=source, revision=revision,
            allow_patterns=["config.json", "model.safetensors", "decision_threshold.json",
                            "tokenizer.json", "tokenizer_config.json", "base_model_reference.json"],
        ))
    reference = json.loads((folder / "base_model_reference.json").read_text())
    decision = json.loads((folder / "decision_threshold.json").read_text())
    tokenizer = AutoTokenizer.from_pretrained(folder, use_fast=True, local_files_only=True)
    tokenizer.model_max_length = reference["max_length"]
    tokenizer.padding_side = "right"
    tokenizer.truncation_side = "right"
    if tokenizer.pad_token_id != reference["pad_token_id"]:
        raise ValueError("Padding token differs from the original evaluated tokenizer")
    model = AutoModelForSequenceClassification.from_pretrained(
        folder, local_files_only=True, use_safetensors=True, dtype=torch.float32,
    ).eval()
    if model.config.model_type != "xlm-roberta" or model.config.num_labels != 2:
        raise ValueError("Expected the original binary XLM-R classifier")
    model.requires_grad_(False)
    device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
    if precision not in {"auto", "fp32", "fp16_mixed"}:
        raise ValueError("Precision must be auto, fp32, or fp16_mixed")
    mixed = device.type == "cuda" if precision == "auto" else precision == "fp16_mixed"
    if mixed and device.type != "cuda":
        raise ValueError("FP16 mixed inference requires a CUDA device; use fp32 on CPU")
    model.to(device)
    return model, tokenizer, reference, decision, mixed


def classify_texts(texts, classifier):
    """Classify complete raw requests with context; no chat template is applied."""
    import torch
    if isinstance(texts, str):
        texts = [texts]
    if not texts or any(not isinstance(text, str) for text in texts):
        raise ValueError("Supply one or more strings")
    model, tokenizer, reference, decision, mixed = classifier
    encoded = tokenizer(
        texts, truncation=True, max_length=reference["max_length"], padding=True,
        pad_to_multiple_of=8, return_token_type_ids=False, return_tensors="pt",
    )
    device = next(model.parameters()).device
    encoded = {key: value.to(device) for key, value in encoded.items()}
    amp = torch.autocast(device_type="cuda", dtype=torch.float16) if mixed else contextlib.nullcontext()
    with torch.inference_mode(), amp:
        logits = model(**encoded).logits
    # The archived worker computes softmax after converting logits to float64.
    values = torch.softmax(logits.double(), dim=-1)[:, decision["positive_class_id"]].cpu().tolist()
    threshold = float(decision["threshold"])
    return [{"decision": "block" if p >= threshold else "allow",
             "block_probability": p, "threshold": threshold} for p in values]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--text", help="Complete raw request and its visible context")
    group.add_argument("--text-file", type=Path, help="UTF-8 file containing one complete request")
    parser.add_argument("--source", default=DEFAULT_REPO, help="Hub repository or local folder")
    parser.add_argument("--revision", help="Optional exact release commit")
    parser.add_argument("--device", help="cpu, cuda, or cuda:0; default selects CUDA if available")
    parser.add_argument("--precision", choices=["auto", "fp32", "fp16_mixed"], default="auto")
    args = parser.parse_args()
    text = args.text if args.text is not None else args.text_file.read_text(encoding="utf-8")
    classifier = load_classifier(args.source, args.revision, args.device, args.precision)
    print(json.dumps(classify_texts(text, classifier)[0], ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()