from __future__ import annotations
from pathlib import Path
import torch

from .models.factory import build_model
from .utils.paths import ARTIFACTS_DIR
from .utils.logging import log

def main() -> None:
    device = torch.device("cpu")
    model = build_model("resnet18", num_classes=10, pretrained=False).to(device)
    ckpt = ARTIFACTS_DIR / "checkpoints" / "best.pt"
    state = torch.load(ckpt, map_location=device)
    model.load_state_dict(state["model"])
    model.eval()

    dummy = torch.randn(1, 3, 32, 32, device=device)
    onnx_path = ARTIFACTS_DIR / "model.onnx"
    torch.onnx.export(
        model, dummy, onnx_path.as_posix(),
        input_names=["input"], output_names=["logits"],
        opset_version=17, dynamic_axes=None
    )
    log(f"Exported ONNX → {onnx_path}")

if __name__ == "__main__":
    main()
