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

from typing import Any, Optional, Sequence

import torch
import torch.nn.functional as F
from PIL import Image
from utils import get_device
from utils.image import crop_image_by_bbox, extract_masked_region

from .filter import MaskScoringConfig, select_best_mask
from .model_loader import (
    get_dino_dim,
    load_dino_model,
    load_hash_compressor,
    load_owlv2_model,
    load_sam_model,
)
from .proposal import (
    detect_objects_batch,
    generate_proposals_batch,
)
from .proposal.core import DetectionResult


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(
        owlv2_model=getattr(
            config.model, "owlv2_model", "google/owlv2-base-patch16-ensemble"
        ),
        dino_model=config.model.dino_model,
        sam_model=config.model.sam_model,
        hash_bits=config.model.compression_dim,
        compressor_path=config.model.compressor_path,
        mask_scoring_config=getattr(config.model, "mask_scoring_config", None),
        score_threshold=getattr(config.model, "score_threshold", 0.25),
        postprocess_threshold=getattr(config.model, "postprocess_threshold", 0.1),
    )


class HashPipeline:
    """Pipeline for OWLv2 detection + SAM segmentation + DINO features + Hash compression.

    Pipeline flow:
        Images + Text Labels -> OWLv2 (detections) -> SAM (masks) -> Filter (best mask) ->
        Crop (OWLv2 box) -> DINO (features) -> Hash (binary codes)

    Example:
        pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512)
        images = [Image.open("path/to/image.jpg")]
        text_labels = ["object"]
        hash_bits = pipeline.forward_batch(images, text_labels)  # Returns [N, 512]
    """

    def __init__(
        self,
        owlv2_model: str = "google/owlv2-base-patch16-ensemble",
        dino_model: str = "facebook/dinov2-large",
        sam_model: str = "facebook/sam2.1-hiera-large",
        hash_bits: int = 512,
        compressor_path: Optional[str] = None,
        mask_scoring_config: Optional["MaskScoringConfig"] = None,
        score_threshold: float = 0.25,
        postprocess_threshold: float = 0.1,
    ):
        super().__init__()

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

        # OWLv2 detection settings.
        self.owlv2_processor, self.owlv2_model = load_owlv2_model(
            model_name=owlv2_model
        )
        self.score_threshold = score_threshold
        self.postprocess_threshold = postprocess_threshold

        # Mask scoring config for filter step.
        self.mask_scoring_config = mask_scoring_config

        # SAM2 model for segmentation.
        self.sam_processor, self.sam_model = load_sam_model(model_name=sam_model)

        # DINO model for feature extraction.
        self.dino_processor, self.dino = load_dino_model(model_name=dino_model)
        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 detect_batch(
        self,
        images: Sequence[Image.Image],
        text_labels: list[str],
    ) -> list[list[DetectionResult]]:
        """Detect objects in a batch of images using OWLv2.

        Args:
            images: Sequence of PIL Images.
            text_labels: Text labels for all images (same labels used for each image).

        Returns:
            List of lists of DetectionResult dicts, one inner list per image.
        """
        image_list = list(images)
        if not image_list:
            return []

        text_labels_per_image = [text_labels] * len(image_list)
        return detect_objects_batch(
            self.owlv2_model,
            self.owlv2_processor,
            image_list,
            text_labels_per_image,
            score_threshold=self.score_threshold,
            postprocess_threshold=self.postprocess_threshold,
        )

    def segment_batch(
        self,
        images: Sequence[Image.Image],
        bboxes_per_image: list[list[list[float]]],
    ) -> list[list[dict[str, Any]]]:
        """Segment objects in images using SAM2 with bounding box prompts.

        Args:
            images: Sequence of PIL Images.
            bboxes_per_image: Bounding boxes per image as [[[x1,y1,x2,y2], ...], ...].

        Returns:
            List of lists of mask dictionaries, one inner list per image.
        """
        image_list = list(images)
        if not image_list:
            return []

        return generate_proposals_batch(
            self.sam_model,
            self.sam_processor,
            image_list,
            bboxes_per_image,
        )

    def filter_batch(
        self,
        images: Sequence[Image.Image],
        masks_per_image: list[list[dict[str, Any]]],
    ) -> list[Image.Image]:
        """Filter masks and extract best masked region for each image.

        Args:
            images: Sequence of PIL Images.
            masks_per_image: Masks per image from segment_batch.

        Returns:
            List of PIL Images, one per input image (original if no valid masks).
        """
        image_list = list(images)
        if not image_list:
            return []

        filtered_images: list[Image.Image] = []
        for index, image in enumerate(image_list):
            masks = masks_per_image[index] if index < len(masks_per_image) else []
            if not masks:
                filtered_images.append(image)
                continue

            best_mask = select_best_mask(
                masks,
                image_shape=(image.height, image.width),
                config=self.mask_scoring_config,
            )
            if best_mask is None:
                filtered_images.append(image)
                continue

            filtered_images.append(extract_masked_region(image, best_mask["segment"]))

        return filtered_images

    def crop_batch(
        self,
        images: Sequence[Image.Image],
        masks_per_image: list[list[dict[str, Any]]],
        detections_per_image: list[list[DetectionResult]],
    ) -> list[Image.Image]:
        """Crop filtered images using OWLv2 detection boxes.

        Args:
            images: Sequence of PIL Images after filter_batch.
            masks_per_image: Masks per image from segment_batch.
            detections_per_image: Detection results per image from detect_batch.

        Returns:
            List of cropped PIL Images. Returns original image when no detection exists.
        """
        image_list = list(images)
        if not image_list:
            return []

        cropped_images: list[Image.Image] = []
        for index, image in enumerate(image_list):
            detections = (
                detections_per_image[index] if index < len(detections_per_image) else []
            )
            if detections:
                best_detection = max(detections, key=lambda d: d["score"])
                cropped_images.append(crop_image_by_bbox(image, best_detection["bbox"]))
                continue

            cropped_images.append(image)

        return cropped_images

    def extract_dino_batch(self, images: Sequence[Image.Image]) -> torch.Tensor:
        """Extract DINO tokens from a batch of images.

        Args:
            images: Sequence of PIL Images.

        Returns:
            Last hidden state tokens of shape [B, N, dim].
        """
        image_list = list(images)
        if not image_list:
            return torch.empty(
                (0, 1, self.dino_dim), dtype=torch.float32, device=self.device
            )

        inputs = self.dino_processor(images=image_list, return_tensors="pt").to(
            self.device
        )

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

    def compress_batch(self, tokens: torch.Tensor) -> torch.Tensor:
        """Compress DINO tokens to binary hash codes.

        Args:
            tokens: DINO tokens of shape [B, N, dim].

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

    def forward_batch(
        self,
        images: Sequence[Image.Image],
        text_labels: list[str],
        batch_size: int = 32,
    ) -> torch.Tensor:
        """Process a batch of images through the full pipeline.

        Args:
            images: Sequence of PIL Images.
            text_labels: Text labels for detection (same for all images).
            batch_size: Batch size for DINO feature extraction.

        Returns:
            Binary hash codes of shape [N, hash_bits] as int32.
        """
        if batch_size <= 0:
            raise ValueError("batch_size must be greater than 0")

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

        detections = self.detect_batch(image_list, text_labels)
        bboxes = [[d["bbox"] for d in dets] for dets in detections]
        masks = self.segment_batch(image_list, bboxes)
        processed = self.filter_batch(image_list, masks)
        processed = self.crop_batch(processed, masks, detections)

        all_bits: list[torch.Tensor] = []
        for i in range(0, len(processed), batch_size):
            sub_batch = processed[i : i + batch_size]
            tokens = self.extract_dino_batch(sub_batch)
            bits = self.compress_batch(tokens)
            all_bits.append(bits)

        return torch.cat(all_bits, dim=0)

    def extract_features_dataset(
        self,
        images: Sequence[Image.Image],
        text_labels: list[str],
        batch_size: int = 32,
    ) -> torch.Tensor:
        """Extract normalized DINO features from a batch of images.

        Args:
            images: Sequence of PIL Images.
            text_labels: Text labels for detection (same for all images).
            batch_size: Batch size for DINO feature extraction.

        Returns:
            Normalized DINO features of shape [N, dino_dim].
        """
        if batch_size <= 0:
            raise ValueError("batch_size must be greater than 0")

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

        detections = self.detect_batch(image_list, text_labels)
        bboxes = [[d["bbox"] for d in dets] for dets in detections]
        masks = self.segment_batch(image_list, bboxes)
        processed = self.filter_batch(image_list, masks)
        processed = self.crop_batch(processed, masks, detections)

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

        return torch.cat(all_features, dim=0)