"""Convert a HuggingFace NSFW classifier to ONNX for fast CPU inference.

Default model: LukeJacob2023/nsfw-image-detector — Vision Transformer, 5 classes
(drawings, hentai, neutral, porn, sexy), 93% test accuracy.

Override via --model to use a different HF repo with the same 5-class schema.

Run this only inside the Docker build stage; do NOT keep transformers/torch in
the runtime image.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

DEFAULT_MODEL = "LukeJacob2023/nsfw-image-detector"


def convert(model_id: str, out_path: Path, cache_dir: Path) -> None:
    # Lazy imports — runtime image never loads these
    import torch  # type: ignore
    from transformers import AutoImageProcessor, AutoModelForImageClassification  # type: ignore

    cache_dir.mkdir(parents=True, exist_ok=True)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"Loading {model_id} from HuggingFace…", flush=True)
    processor = AutoImageProcessor.from_pretrained(model_id, cache_dir=str(cache_dir))
    model = AutoModelForImageClassification.from_pretrained(
        model_id, cache_dir=str(cache_dir)
    )
    model.eval()

    # Persist the label order — runtime needs to know which index = which class
    id2label = model.config.id2label
    labels = [id2label[i] for i in sorted(id2label.keys(), key=int)]
    print(f"Model labels: {labels}", flush=True)

    expected = ["drawings", "hentai", "neutral", "porn", "sexy"]
    if sorted(labels) != sorted(expected):
        raise RuntimeError(
            f"Model {model_id} returns labels {labels}, "
            f"expected {expected}. Use a different model or update CATEGORIES."
        )

    # ImageProcessor settings — image_mean, image_std, size
    size = processor.size
    if isinstance(size, dict):
        height = size.get("height") or size.get("shortest_edge") or 224
        width = size.get("width") or size.get("shortest_edge") or 224
    else:
        height = width = int(size)

    metadata = {
        "model_id": model_id,
        "labels": labels,
        "image_size": [int(height), int(width)],
        "image_mean": [float(x) for x in processor.image_mean],
        "image_std": [float(x) for x in processor.image_std],
        "rescale_factor": float(getattr(processor, "rescale_factor", 1.0 / 255.0)),
    }

    # Export to ONNX
    print(f"Exporting to ONNX → {out_path}", flush=True)
    dummy = torch.randn(1, 3, height, width)

    torch.onnx.export(
        model,
        dummy,
        str(out_path),
        input_names=["pixel_values"],
        output_names=["logits"],
        dynamic_axes={"pixel_values": {0: "batch"}, "logits": {0: "batch"}},
        opset_version=14,
        do_constant_folding=True,
    )

    # Save metadata next to the model
    meta_path = out_path.with_suffix(".json")
    meta_path.write_text(json.dumps(metadata, indent=2))

    size_mb = out_path.stat().st_size / (1024 * 1024)
    print(f"Done. ONNX size: {size_mb:.1f} MB", flush=True)
    print(f"Metadata written to {meta_path}", flush=True)


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model",
        default=os.getenv("NSFW_MODEL_ID", DEFAULT_MODEL),
        help=f"HuggingFace model ID (default: {DEFAULT_MODEL})",
    )
    parser.add_argument("--out", default="/app/model/nsfw.onnx")
    parser.add_argument("--cache", default="/tmp/nsfw-build/hf-cache")
    args = parser.parse_args(argv)

    convert(args.model, Path(args.out), Path(args.cache))
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))