"""SAM + DINO + Hash compression pipeline."""

from typing import Optional, Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from utils import get_device

from .filter import select_best_mask
from .model_loader import (
    get_dino_dim,
    load_dino_model,
    load_hash_compressor,
    load_sam_model,
)
from .proposal import (
    extract_masked_region,
    generate_proposals,
    generate_proposals_batch,
)


def create_pipeline_from_config(config) -> "HashPipeline":
    """Create HashPipeline from a config object.

    Args:
        config: Configuration object with model settings.

    Returns:
        Initialized HashPipeline.
    """
    return HashPipeline(
        dino_model=config.model.dino_model,
        sam_model=config.model.sam_model,
        sam_min_mask_area=config.model.sam_min_mask_area,
        sam_max_masks=config.model.sam_max_masks,
        hash_bits=config.model.compression_dim,
        compressor_path=config.model.compressor_path,
    )


class HashPipeline(nn.Module):
    """Pipeline for SAM segmentation + DINO features + Hash compression.

    Pipeline flow:
        PIL Image -> SAM (largest object mask) -> DINO (features) -> Hash (binary codes)

    Example:
        pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512)
        image = Image.open("path/to/image.jpg")
        hash_bits = pipeline(image)  # Returns [1, 512] binary bits
    """

    def __init__(
        self,
        dino_model: str = "facebook/dinov2-large",
        sam_model: str = "facebook/sam2.1-hiera-large",
        sam_min_mask_area: int = 100,
        sam_max_masks: int = 10,
        hash_bits: int = 512,
        compressor_path: Optional[str] = None,
    ):
        super().__init__()

        # Device for model placement.
        self.device = get_device()

        # SAM2 filter settings.
        self.sam_min_mask_area = sam_min_mask_area
        self.sam_max_masks = sam_max_masks

        # Load models.
        self.sam_processor, self.sam_model = load_sam_model(model_name=sam_model)
        self.dino_processor, self.dino = load_dino_model(model_name=dino_model)

        # DINO feature dimension based on model size.
        self.dino_dim = get_dino_dim(dino_model)

        # Hash compressor for binarizing DINO features.
        self.hash_compressor = load_hash_compressor(
            input_dim=self.dino_dim,
            hash_bits=hash_bits,
            compressor_path=compressor_path,
        )

    @property
    def hash_bits(self) -> int:
        """Number of bits in the hash code."""
        return self.hash_compressor.hash_bits

    def _segment_with_sam(
        self, image: Image.Image, bboxes: list[list[float]]
    ) -> Image.Image:
        """Segment image with SAM and extract the best object mask.

        Args:
            image: Input PIL Image.
            bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].

        Returns:
            Masked image containing only the best object, or original if no masks.
        """
        masks = generate_proposals(
            self.sam_model,
            self.sam_processor,
            image,
            bboxes,
        )

        masks = _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks)

        if not masks:
            return image

        best_mask = select_best_mask(masks, image_shape=(image.height, image.width))
        if best_mask is None:
            return image
        return extract_masked_region(image, best_mask["segment"])

    def _segment_with_sam_dataset(
        self,
        images: Sequence[Image.Image],
        bboxes_per_image: list[list[list[float]]],
    ) -> list[Image.Image]:
        image_list = list(images)
        masks_dataset = generate_proposals_batch(
            self.sam_model,
            self.sam_processor,
            image_list,
            bboxes_per_image,
        )
        masks_dataset = [
            _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks)
            for masks in masks_dataset
        ]
        selected_images: list[Image.Image] = []
        for image, masks in zip(image_list, masks_dataset):
            if not masks:
                selected_images.append(image)
                continue

            best_mask = select_best_mask(masks, image_shape=(image.height, image.width))
            if best_mask is None:
                selected_images.append(image)
                continue
            selected_images.append(extract_masked_region(image, best_mask["segment"]))

        return selected_images

    def _dino_forward(self, image: Image.Image) -> torch.Tensor:
        """Extract DINO tokens from an image.

        Args:
            image: Input PIL Image.

        Returns:
            Last hidden state tokens of shape [1, N, dim].
        """
        inputs = self.dino_processor(image, return_tensors="pt").to(self.device)

        with torch.no_grad():
            outputs = self.dino(**inputs)
            return outputs.last_hidden_state

    def _dino_forward_batch(self, images: Sequence[Image.Image]) -> torch.Tensor:
        inputs = self.dino_processor(images=list(images), return_tensors="pt").to(
            self.device
        )

        with torch.no_grad():
            outputs = self.dino(**inputs)
            return outputs.last_hidden_state

    def forward(self, image: Image.Image, bboxes: list[list[float]]) -> torch.Tensor:
        """Process a single image through the full pipeline.

        Args:
            image: Input PIL Image.
            bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].

        Returns:
            Binary hash codes of shape [1, hash_bits] as int32.
        """
        image = self._segment_with_sam(image, bboxes)
        tokens = self._dino_forward(image)
        _, _, bits = self.hash_compressor(tokens)
        return bits

    def forward_dataset(
        self,
        images: Sequence[Image.Image],
        bboxes_per_image: list[list[list[float]]],
        batch_size: int = 32,
        apply_sam: bool = True,
    ) -> torch.Tensor:
        if batch_size <= 0:
            raise ValueError("batch_size must be greater than 0")

        image_list = list(images)

        if len(image_list) == 0:
            return torch.empty(
                (0, self.hash_bits), dtype=torch.int32, device=self.device
            )

        if apply_sam:
            processed_images = self._segment_with_sam_dataset(
                image_list, bboxes_per_image
            )
        else:
            processed_images = image_list

        batch_bits: list[torch.Tensor] = []
        for i in range(0, len(processed_images), batch_size):
            batch_images = processed_images[i : i + batch_size]
            tokens = self._dino_forward_batch(batch_images)
            _, _, bits = self.hash_compressor(tokens)
            batch_bits.append(bits)

        return torch.cat(batch_bits, dim=0)

    def extract_features(
        self, image: Image.Image, bboxes: list[list[float]]
    ) -> torch.Tensor:
        """Extract normalized DINO features from an image.

        Args:
            image: Input PIL Image.
            bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].

        Returns:
            Normalized DINO features of shape [1, dino_dim].
        """
        image = self._segment_with_sam(image, bboxes)
        tokens = self._dino_forward(image)
        features = tokens.mean(dim=1)
        return F.normalize(features, dim=-1)

    def extract_features_dataset(
        self,
        images: Sequence[Image.Image],
        bboxes_per_image: list[list[list[float]]],
        batch_size: int = 32,
        apply_sam: bool = True,
    ) -> torch.Tensor:
        if batch_size <= 0:
            raise ValueError("batch_size must be greater than 0")

        image_list = list(images)

        if len(image_list) == 0:
            return torch.empty(
                (0, self.dino_dim), dtype=torch.float32, device=self.device
            )

        if apply_sam:
            processed_images = self._segment_with_sam_dataset(
                image_list, bboxes_per_image
            )
        else:
            processed_images = image_list

        all_features: list[torch.Tensor] = []
        for i in range(0, len(processed_images), batch_size):
            batch_images = processed_images[i : i + batch_size]
            tokens = self._dino_forward_batch(batch_images)
            features = tokens.mean(dim=1)
            all_features.append(F.normalize(features, dim=-1))

        return torch.cat(all_features, dim=0)


def _filter_masks(
    masks: list[dict],
    min_area: int,
    max_masks: int,
) -> list[dict]:
    """Filter masks by area and keep top-N largest."""
    filtered = [m for m in masks if int(m["area"]) >= min_area]
    if not filtered:
        return []
    sorted_masks = sorted(filtered, key=lambda m: m["area"], reverse=True)
    return sorted_masks[:max_masks]