from __future__ import annotations

import argparse
from pathlib import Path


def main() -> int:
    parser = argparse.ArgumentParser(description="Export trained RF-DETR to ONNX for deployment.")
    parser.add_argument("--weights", type=Path, required=True, help="Path to trained checkpoint")
    parser.add_argument("--output-onnx", type=Path, default=Path("model.onnx"), help="Output ONNX file")
    args = parser.parse_args()

    if not args.weights.exists():
        raise SystemExit(f"Weights not found: {args.weights}")

    from rfdetr import RFDETRBase

    model = RFDETRBase(pretrain_weights=str(args.weights))

    # Export to ONNX. This saves to the current directory by default, but we can specify.
    # RF-DETR's export() method saves to 'output/model.onnx' I think, but let's check docs.
    # From earlier fetch: model.export() saves to 'output' dir.
    # But to make it flexible, perhaps run it and then move.

    model.export()  # This should create output/model.onnx

    # Move to desired location
    onnx_path = Path("output/model.onnx")
    if onnx_path.exists():
        onnx_path.rename(args.output_onnx)
        print(f"Exported ONNX: {args.output_onnx}")
    else:
        raise SystemExit("ONNX export failed - check output/ dir")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())