import os
import urllib.request
import torch
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator

def _download_sam_checkpoint(checkpoint_path="outputs/models/sam_vit_b_01ec64.pth"):
    """
    Downloads the Segment Anything Model (SAM) ViT-B checkpoint (350MB).

    Args:
        checkpoint_path (str): Local path to save/check the .pth file.
    """
    url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
    if not os.path.exists(checkpoint_path):
        print(f"Downloading SAM checkpoint to {checkpoint_path}...")
        os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
        urllib.request.urlretrieve(url, checkpoint_path)
        print("Download complete.")
    else:
        print("SAM checkpoint already exists.")

def get_sam_generator(model_type="vit_b", checkpoint_path="outputs/models/sam_vit_b_01ec64.pth", device="cuda"):
    """
    Initializes the SAM model and returns an automatic mask generator.

    Args:
        model_type (str): The SAM architecture version (e.g., 'vit_h').
        checkpoint_path (str): Path to the model weights.
        device (str): Device to run the model on ('cuda' or 'cpu').

    Returns:
        SamAutomaticMaskGenerator: Initialized mask generator object.
    """
    if device == "cuda" and not torch.cuda.is_available():
        device = "cpu"
        print("CUDA not available. Falling back to CPU for SAM.")

    _download_sam_checkpoint(checkpoint_path)

    sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
    sam.to(device=device)
    mask_generator = SamAutomaticMaskGenerator(sam)
    return mask_generator