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

from typing import TYPE_CHECKING, Any

import torch
from transformers import (
    AutoImageProcessor,
    AutoModel,
    Dinov2Model,
    Owlv2ForObjectDetection,
    Owlv2Processor,
    Sam2Model,
    Sam2Processor,
)
from utils import get_device

if TYPE_CHECKING:
    from compressors.hash_compressor import HashCompressor


def load_sam_model(
    model_name: str = "facebook/sam2.1-hiera-large",
) -> tuple[Sam2Processor, Sam2Model]:
    """Load SAM2 processor and model with frozen parameters."""
    device = get_device()

    processor = Sam2Processor.from_pretrained(model_name)
    model = Sam2Model.from_pretrained(model_name).to(device)
    model.eval()
    for param in model.parameters():
        param.requires_grad = False

    return processor, model


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

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

    return processor, dino


def load_owlv2_model(
    model_name: str = "google/owlv2-base-patch16-ensemble",
) -> tuple[Owlv2Processor, Owlv2ForObjectDetection]:
    device = get_device()

    processor = Owlv2Processor.from_pretrained(model_name)
    model = Owlv2ForObjectDetection.from_pretrained(model_name).to(device)
    model.eval()

    return processor, model


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