from __future__ import annotations

from pathlib import Path
from typing import Optional

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

from sam3.model_builder import build_sam3_image_model


def load_model(
    repo_id: str = "TechieMoon/sam3-libero10-procedural-segmentation",
    filename: str = "model.safetensors",
    device: str = "cuda" if torch.cuda.is_available() else "cpu",
    bpe_path: Optional[str] = None,
    local_path: Optional[str | Path] = None,
):
    """Load the LIBERO-10 fine-tuned SAM3 image model.

    The checkpoint stores a direct SAM3 image-model state dict. Build the
    architecture without downloading a base checkpoint, then load this state.
    """
    model_path = Path(local_path) if local_path is not None else Path(
        hf_hub_download(repo_id=repo_id, filename=filename)
    )
    model = build_sam3_image_model(
        bpe_path=bpe_path,
        checkpoint_path=None,
        load_from_HF=False,
        device="cpu",
        eval_mode=True,
        enable_segmentation=True,
        enable_inst_interactivity=False,
    )
    state = load_file(str(model_path), device="cpu")
    model.load_state_dict(state, strict=True)
    model.to(device)
    model.eval()
    return model


if __name__ == "__main__":
    loaded = load_model(device="cpu")
    n_params = sum(p.numel() for p in loaded.parameters())
    print(f"Loaded SAM3 LIBERO-10 model with {n_params:,} parameters.")