#!/usr/bin/env python3
"""Prepare a pinned Hugging Face binary classifier for VibeGuard."""
from __future__ import annotations

import argparse
import json
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

DEFAULT_MODEL = "wagesj45/toxic-comment-classifier"
DEFAULT_OUTPUT = Path("public/models/toxicity")


def make_parser() -> argparse.ArgumentParser:
    command = argparse.ArgumentParser(description=__doc__)
    command.add_argument("action", choices=("inspect", "prepare", "validate"))
    command.add_argument("--model", default=DEFAULT_MODEL)
    command.add_argument("--revision", help="Immutable HF commit; required by prepare")
    command.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    command.add_argument("--toxic-index", type=int)
    command.add_argument("--non-toxic-index", type=int)
    command.add_argument("--max-length", type=int)
    return command


def load_source(model: str, revision: str | None) -> tuple[Path, dict[str, Any]]:
    try:
        from huggingface_hub import model_info, snapshot_download
    except ImportError as error:
        raise SystemExit("Install tools with: python3 -m pip install -r tools/requirements-model.txt") from error
    kwargs: dict[str, Any] = {"repo_id": model}
    if revision: kwargs["revision"] = revision
    source = Path(snapshot_download(**kwargs))
    info = None
    try:
        info = model_info(model, revision=revision)
    except Exception as error:  # A pinned, complete local snapshot is still usable offline.
        print(f"Warning: Hugging Face metadata unavailable; using local snapshot metadata: {error}", file=sys.stderr)
    card = getattr(info, "cardData", None) or {}
    license_id = card.get("license")
    if not license_id and (source / "LICENSE").exists():
        license_text = (source / "LICENSE").read_text(errors="replace").lower()
        if "apache license" in license_text and "version 2.0" in license_text:
            license_id = "apache-2.0"
    return source, {"source": model, "revision": revision or getattr(info, "sha", "unknown"), "license": license_id}


def inspect_source(source: Path) -> dict[str, Any]:
    config = json.loads((source / "config.json").read_text())
    labels = {str(key): str(value) for key, value in (config.get("id2label") or {}).items()}
    return {
        "architecture": (config.get("architectures") or [config.get("model_type", "unknown")])[0],
        "modelType": config.get("model_type"),
        "numLabels": config.get("num_labels") or len(labels),
        "id2label": labels,
        "maxLength": config.get("max_position_embeddings"),
        "files": sorted(path.name for path in source.iterdir()),
    }


def resolve_labels(details: dict[str, Any], toxic_override: int | None, non_toxic_override: int | None) -> tuple[int, int, dict[str, int]]:
    names = {str(index): value.lower().replace("-", "_").replace(" ", "_") for index, value in details["id2label"].items()}
    indexed = {name: int(index) for index, name in names.items()}
    toxic = toxic_override if toxic_override is not None else next((index for name, index in indexed.items() if name in ("toxic", "toxicity")), None)
    non_toxic = non_toxic_override if non_toxic_override is not None else next((index for name, index in indexed.items() if name in ("non_toxic", "nontoxic", "clean", "not_toxic")), None)
    if toxic is None or non_toxic is None or toxic == non_toxic:
        raise ValueError("Could not resolve distinct labels; pass --toxic-index and --non-toxic-index")
    return toxic, non_toxic, indexed


def prepare(args: argparse.Namespace) -> None:
    if not args.revision: raise SystemExit("prepare requires --revision for reproducible release artifacts")
    source, source_metadata = load_source(args.model, args.revision)
    license_id = str(source_metadata.get("license") or "").lower()
    if license_id != "apache-2.0": raise SystemExit(f"Expected Apache-2.0 model metadata, found license={source_metadata.get('license')!r}")
    details = inspect_source(source)
    if details["numLabels"] != 2: raise SystemExit(f"Expected a binary classifier, found {details['numLabels']}")
    toxic, non_toxic, names = resolve_labels(details, args.toxic_index, args.non_toxic_index)
    max_length = args.max_length or details["maxLength"] or 512
    with tempfile.TemporaryDirectory(prefix="vibeguard-model-") as temporary:
        export_dir = Path(temporary) / "onnx"
        float_model = export_dir / "model.onnx"
        export_onnx(source, float_model, max_length)
        quantized = export_dir / "model_quantized.onnx"
        quantize(float_model, quantized)
        args.output.mkdir(parents=True, exist_ok=True)
        for name in ("tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "vocab.txt", "merges.txt"):
            if (source / name).exists(): shutil.copy2(source / name, args.output / name)
        shutil.copy2(source / "config.json", args.output / "config.json")
        (args.output / "onnx").mkdir(exist_ok=True)
        shutil.copy2(quantized, args.output / "onnx" / "model_quantized.onnx")
        manifest = {"source": args.model, "revision": source_metadata["revision"], "architecture": details["architecture"], "labels": {"toxic": toxic, "nonToxic": non_toxic, "names": names}, "maxLength": max_length, "quantization": "int8-dynamic", "runtime": "onnxruntime-web-wasm", "license": "Apache-2.0"}
        (args.output / "model-manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    print(json.dumps(manifest, indent=2))


def quantize(source: Path, destination: Path) -> None:
    try:
        from onnxruntime.quantization import QuantType, quantize_dynamic
    except ImportError as error:
        raise SystemExit("Install tools with: python3 -m pip install -r tools/requirements-model.txt") from error
    quantize_dynamic(str(source), str(destination), weight_type=QuantType.QInt8, per_channel=True, reduce_range=True)


def export_onnx(source: Path, destination: Path, max_length: int) -> None:
    try:
        import torch
        from transformers import AutoModelForSequenceClassification, AutoTokenizer
    except ImportError as error:
        raise SystemExit("Install model tooling with: python3 -m pip install -r tools/requirements-model.txt") from error

    tokenizer = AutoTokenizer.from_pretrained(str(source), local_files_only=True)
    model = AutoModelForSequenceClassification.from_pretrained(str(source), local_files_only=True)
    model.eval()
    encoded = tokenizer("VibeGuard export calibration text.", return_tensors="pt", truncation=True, max_length=min(max_length, 32), padding="max_length")

    class ClassifierWrapper(torch.nn.Module):
        def __init__(self, wrapped: torch.nn.Module) -> None:
            super().__init__()
            self.wrapped = wrapped

        def forward(self, input_ids: Any, attention_mask: Any) -> Any:
            return self.wrapped(input_ids=input_ids, attention_mask=attention_mask).logits

    destination.parent.mkdir(parents=True, exist_ok=True)
    with torch.no_grad():
        torch.onnx.export(
            ClassifierWrapper(model),
            (encoded["input_ids"], encoded["attention_mask"]),
            str(destination),
            input_names=["input_ids", "attention_mask"],
            output_names=["logits"],
            dynamic_axes={"input_ids": {0: "batch", 1: "sequence"}, "attention_mask": {0: "batch", 1: "sequence"}, "logits": {0: "batch"}},
            opset_version=17,
            do_constant_folding=True,
            dynamo=False,
        )


def validate(args: argparse.Namespace) -> None:
    manifest_path = args.output / "model-manifest.json"
    model_path = args.output / "onnx" / "model_quantized.onnx"
    if not manifest_path.is_file() or not model_path.is_file(): raise SystemExit(f"Missing generated model files under {args.output}")
    manifest = json.loads(manifest_path.read_text())
    if manifest.get("quantization") != "int8-dynamic" or str(manifest.get("license")).lower() != "apache-2.0": raise SystemExit("Unexpected model quantization or license")
    try:
        import onnx
        onnx.checker.check_model(str(model_path))
    except ImportError as error:
        raise SystemExit("Install tools with: python3 -m pip install -r tools/requirements-model.txt") from error
    print(f"Validated {model_path} ({model_path.stat().st_size / 1024 / 1024:.1f} MiB)")


def run(command: list[str]) -> None:
    print("+", " ".join(command))
    subprocess.run(command, check=True)


def main() -> None:
    args = make_parser().parse_args()
    if args.action == "inspect":
        source, metadata = load_source(args.model, args.revision)
        print(json.dumps({**metadata, **inspect_source(source)}, indent=2))
    elif args.action == "prepare": prepare(args)
    else: validate(args)


if __name__ == "__main__": main()
