#!/usr/bin/env python3
# Author: Ilya Zelenskiy (proxy3d)
# Telegram channel: https://t.me/greenruff
from __future__ import annotations

import argparse
import json
from pathlib import Path

ROOT = Path(__file__).resolve().parent
EXPECTED_WEIGHT_SIZE = 1_112_284_984
REQUIRED = [
    "README.md",
    "LICENSE",
    "config.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "special_tokens_map.json",
    "thresholds.json",
    "benchmarks.json",
    "model.safetensors",
]


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--load-model", action="store_true", help="Also load tokenizer/model and run one forward pass")
    args = p.parse_args()

    missing = [name for name in REQUIRED if not (ROOT / name).is_file()]
    if missing:
        raise SystemExit("Missing required files: " + ", ".join(missing))

    weight_size = (ROOT / "model.safetensors").stat().st_size
    if weight_size != EXPECTED_WEIGHT_SIZE:
        raise SystemExit(
            f"Wrong model.safetensors size: {weight_size:,} bytes; expected {EXPECTED_WEIGHT_SIZE:,}. "
            "Copy the weight from the seed456 Drive folder."
        )

    config = json.loads((ROOT / "config.json").read_text(encoding="utf-8"))
    thresholds = json.loads((ROOT / "thresholds.json").read_text(encoding="utf-8"))
    id2label = config.get("id2label", {})
    labels = [id2label[str(i)] for i in range(len(id2label))]

    if len(labels) != 28:
        raise SystemExit(f"config.json has {len(labels)} labels; expected 28")
    if config.get("problem_type") != "multi_label_classification":
        raise SystemExit("config.json problem_type is not multi_label_classification")
    if set(labels) != set(thresholds):
        raise SystemExit("thresholds.json label set does not match config.json")

    print("Static release checks: OK")
    print(f"model.safetensors: {weight_size:,} bytes")
    print("labels: 28")
    print("problem_type: multi_label_classification")

    if args.load_model:
        try:
            import torch
            from transformers import AutoModelForSequenceClassification, AutoTokenizer
        except ImportError as exc:
            raise SystemExit("Install torch + transformers first: pip install -r requirements.txt") from exc

        tokenizer = AutoTokenizer.from_pretrained(ROOT, use_fast=True, local_files_only=True)
        model = AutoModelForSequenceClassification.from_pretrained(ROOT, local_files_only=True).eval()
        encoded = tokenizer(
            "Я рад, что всё получилось.",
            return_tensors="pt",
            truncation=True,
            max_length=64,
        )
        with torch.inference_mode():
            logits = model(**encoded).logits
        if tuple(logits.shape) != (1, 28):
            raise SystemExit(f"Unexpected logits shape: {tuple(logits.shape)}")
        print("Tokenizer/model load + forward pass: OK")


if __name__ == "__main__":
    main()