"""
Export the OpenCLIP ViT-B/32 visual encoder to ONNX.

Run once on any machine with Python + pip (no GPU needed):

    pip install open-clip-torch onnx
    python -m app.services.vision.export_models [--models-dir /data/models]

Produces:
    embed/visual.onnx   (~350 MB)
"""
import argparse
import logging
from pathlib import Path

logger = logging.getLogger(__name__)


def export_openclip_visual(models_dir: Path):
    import torch
    import open_clip

    out_dir = models_dir / "embed"
    out_dir.mkdir(parents=True, exist_ok=True)
    visual_path = out_dir / "visual.onnx"
    if visual_path.exists():
        logger.info("OpenCLIP visual.onnx already exists, skipping export")
        return

    logger.info("Loading OpenCLIP ViT-B-32 laion2b_s34b_b79k...")
    model, _, _ = open_clip.create_model_and_transforms(
        "ViT-B-32", pretrained="laion2b_s34b_b79k"
    )
    model.eval()

    logger.info("Exporting visual encoder → %s", visual_path)
    dummy = torch.randn(1, 3, 224, 224)
    torch.onnx.export(
        model.visual,
        dummy,
        str(visual_path),
        input_names=["image"],
        output_names=["embedding"],
        dynamic_axes={"image": {0: "batch"}},
        opset_version=14,
        dynamo=False,
    )
    logger.info("Visual encoder exported (%.1f MB)", visual_path.stat().st_size / 1e6)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--models-dir", type=Path, default=Path("/data/models"))
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    args.models_dir.mkdir(parents=True, exist_ok=True)
    export_openclip_visual(args.models_dir)


if __name__ == "__main__":
    main()