#!/usr/bin/env python3
"""Export conservationxlabs/miewid-msv3 to ONNX.

Clones the HF model, loads via trust_remote_code, traces forward(),
and exports as ONNX. Verifies output with onnxruntime.

The exported model accepts `[B, 3, 440, 440]` float32 and returns
`[B, 2152]` float32 embedding vectors (BatchNorm head output).

Usage::

    python scripts/export.py [--upload]
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
HF_REPO = "james-burgess/miewid"
LOCAL_DIR = Path("/tmp/miewid-msv3")


def clone_repo() -> Path:
    import subprocess

    if LOCAL_DIR.exists():
        subprocess.run(["git", "-C", str(LOCAL_DIR), "pull", "--ff-only"], check=False)
    else:
        subprocess.run(
            [
                "git",
                "clone",
                "--depth",
                "1",
                "https://huggingface.co/conservationxlabs/miewid-msv3",
                str(LOCAL_DIR),
            ],
            check=True,
        )
    return LOCAL_DIR


def export() -> tuple[Path, int]:
    import numpy as np
    import torch

    repo = clone_repo()

    sys.path.insert(0, str(repo))

    from transformers import AutoModel

    print("Loading model from local clone...")
    model = AutoModel.from_pretrained(
        str(repo), trust_remote_code=True, torch_dtype=torch.float32
    )
    model.eval()
    print(f"Model loaded: {type(model).__name__}")

    # Dry-run to get output dim
    with torch.inference_mode():
        dummy = torch.randn(1, 3, 440, 440)
        result = model(dummy)
        out_dim = result.shape[-1]
        print(f"Dry run: {dummy.shape} -> {result.shape} (dim={out_dim})")

    # Export
    dest = REPO_ROOT / "miewid.onnx"

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

    size_mb = dest.stat().st_size / 1_048_576
    print(f"Exported -> {dest} ({size_mb:.1f} MB)")

    _verify(dest, out_dim)
    return dest, out_dim


def _verify(onnx_path: Path, expected_dim: int) -> None:
    import numpy as np
    import onnxruntime as ort

    session = ort.InferenceSession(str(onnx_path))
    inp_name = session.get_inputs()[0].name
    out_name = session.get_outputs()[0].name

    dummy = np.random.randn(1, 3, 440, 440).astype(np.float32)
    result = session.run([out_name], {inp_name: dummy})[0]

    assert result.shape[-1] == expected_dim, f"{result.shape[-1]} != {expected_dim}"
    print(f"  Verified: {dummy.shape} -> {result.shape}")


def upload(filepath: Path, out_dim: int) -> None:
    try:
        from huggingface_hub import HfApi, create_repo
    except ImportError:
        print("Install huggingface_hub: pip install huggingface_hub")
        return

    api = HfApi()
    create_repo(HF_REPO, repo_type="model", exist_ok=True)

    size_mb = filepath.stat().st_size / 1_048_576

    # Upload the ONNX model
    api.upload_file(
        path_or_fileobj=str(filepath),
        path_in_repo="miewid.onnx",
        repo_id=HF_REPO,
        commit_message=f"miewid.onnx ({size_mb:.1f} MB, dim={out_dim})",
    )

    # Upload metadata
    for meta_file in ("README.md", "config.json"):
        meta_path = REPO_ROOT / meta_file
        if meta_path.exists():
            api.upload_file(
                path_or_fileobj=str(meta_path),
                path_in_repo=meta_file,
                repo_id=HF_REPO,
                commit_message=f"Update {meta_file}",
            )
            print(f"  Uploaded {meta_file}")

    print(f"Uploaded -> https://huggingface.co/{HF_REPO}")


def main() -> None:
    parser = argparse.ArgumentParser(description="Export MiewID HF model to ONNX")
    parser.add_argument(
        "--upload",
        action="store_true",
        help="Upload model + metadata to HF after export",
    )
    args = parser.parse_args()

    onnx_path, out_dim = export()
    if args.upload:
        upload(onnx_path, out_dim)


if __name__ == "__main__":
    main()