"""Model loading utilities for DINO, SAM2 and HashCompressor."""

from typing import TYPE_CHECKING


import torch
from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGenerationPipeline

from .common import get_device

if TYPE_CHECKING:
    from compressors.hash_compressor import HashCompressor


def load_sam_model(
    model_name: str = "facebook/sam2.1-hiera-large",
) -> MaskGenerationPipeline:
    device = get_device()

    return pipeline(
        task="mask-generation",
        model=model_name,
        device=device,
    )


def load_dino_model(
    model_name: str = "facebook/dinov2-large",
) -> tuple[AutoImageProcessor, AutoModel]:
    device = get_device()

    processor = AutoImageProcessor.from_pretrained(model_name)
    dino = AutoModel.from_pretrained(model_name).to(device)
    dino.eval()

    return processor, dino


def get_dino_dim(model_name: str) -> int:
    if "large" in model_name.lower():
        return 1024
    return 768


def load_hash_compressor(
    input_dim: int = 1024,
    hash_bits: int = 512,
    compressor_path: str | None = None,
) -> "HashCompressor":
    from compressors.hash_compressor import HashCompressor

    device = get_device()
    compressor = HashCompressor(input_dim=input_dim, hash_bits=hash_bits).to(device)

    if compressor_path is not None:
        compressor.load_state_dict(torch.load(compressor_path, map_location=device))
        print(f"[OK] Loaded HashCompressor from {compressor_path}")

    return compressor