"""Load SafeAge Guard as a binary classifier over its pinned Qwen base."""
import argparse
import json
from pathlib import Path

DEFAULT_REPO = "mehedi4/SafeAge-Guard"


def load_guard(source=DEFAULT_REPO, revision=None, device=None):
    import torch
    from huggingface_hub import snapshot_download
    from peft import PeftModel
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    folder = Path(source)
    if not folder.is_dir():
        folder = Path(snapshot_download(
            repo_id=source, revision=revision,
            allow_patterns=["adapter_config.json", "adapter_model.safetensors",
                            "base_model_reference.json", "config.json", "decision_threshold.json",
                            "tokenizer.json", "tokenizer_config.json", "chat_template.jinja"],
        ))
    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.padding_side = "right"
    tokenizer.truncation_side = "right"
    tokenizer.model_max_length = reference["max_length"]
    if tokenizer.pad_token_id != reference["pad_token_id"]:
        raise ValueError("Tokenizer padding token differs from the original checkpoint")
    labels = {int(k): v for k, v in reference["id2label"].items()}
    base = AutoModelForSequenceClassification.from_pretrained(
        reference["model_source"], revision=reference["model_revision"],
        num_labels=reference["num_labels"], id2label=labels,
        label2id={v: k for k, v in labels.items()},
        use_safetensors=True, attn_implementation="sdpa", dtype=torch.float32,
    )
    base.config.pad_token_id = reference["pad_token_id"]
    base.config.problem_type = "single_label_classification"
    base.config.use_cache = False
    model = PeftModel.from_pretrained(
        base, folder, is_trainable=False, local_files_only=True,
    ).eval()
    model.requires_grad_(False)
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    return model, tokenizer, reference, decision


def classify_texts(texts, guard):
    """Classify raw requests; do not apply a chat template."""
    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 = guard
    encoded = tokenizer(
        texts, truncation=True, max_length=reference["max_length"],
        padding=True, pad_to_multiple_of=8, return_tensors="pt",
    )
    device = next(model.parameters()).device
    encoded = {key: value.to(device) for key, value in encoded.items()}
    with torch.inference_mode():
        probabilities = torch.softmax(model(**encoded).logits.float(), dim=-1)
    values = probabilities[:, 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 with the 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 model folder")
    parser.add_argument("--revision", help="Optional Hub commit to pin the adapter release")
    parser.add_argument("--device", help="For example cpu or cuda; default uses CUDA when available")
    args = parser.parse_args()
    text = args.text if args.text is not None else args.text_file.read_text(encoding="utf-8")
    guard = load_guard(source=args.source, revision=args.revision, device=args.device)
    print(json.dumps(classify_texts(text, guard)[0], ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()