from __future__ import annotations

import argparse
from pathlib import Path

import torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file

from repostguard.config import load_config
from repostguard.models import build_model


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--revision", default="v1.0.0")
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()

    root = Path(
        snapshot_download(
            repo_id="LLL640/RepostGuard-Lite-M3-train-v3",
            revision=args.revision,
        )
    )
    config = load_config(root / "resolved_config.yaml")
    model = build_model(config, load_pretrained=False)
    model.load_state_dict(load_file(root / "model.safetensors"), strict=True)
    model = model.to(args.device).eval()

    # Replace this tensor with repository preprocessing for a real RGB image.
    image = torch.zeros(1, 3, 224, 224, device=args.device)
    with torch.inference_mode():
        score = torch.sigmoid(model(image)["logits"])
    print(float(score.item()))


if __name__ == "__main__":
    main()