"""SAM mask proposal generation via bounding box prompts."""

from typing import Any, Sequence, TypedDict

import numpy as np
import torch
from PIL import Image
from transformers import (
    Owlv2ForObjectDetection,
    Owlv2Processor,
    Sam2Model,
    Sam2Processor,
)
from utils import get_device


class DetectionResult(TypedDict):
    bbox: list[float]
    score: float
    label: str


def generate_proposals(
    model: Sam2Model,
    processor: Sam2Processor,
    image: Image.Image,
    bboxes: list[list[float]],
) -> list[dict[str, Any]]:
    """Segment regions in image using SAM2 with bounding box prompts.

    Args:
        model: Sam2Model instance.
        processor: Sam2Processor instance.
        image: PIL Image to segment.
        bboxes: Bounding boxes as [[x1, y1, x2, y2], ...].

    Returns:
        List of mask dictionaries with keys:
            - segment: Binary mask (numpy array)
            - area: Mask area in pixels
            - bbox: Bounding box [x, y, width, height]
            - predicted_iou: Model's confidence in the mask
            - stability_score: Stability score for the mask
    """
    if not bboxes:
        return []

    device = get_device()
    image_rgb = image.convert("RGB")
    input_boxes = [bboxes]
    inputs = processor(
        images=image_rgb,
        input_boxes=input_boxes,
        return_tensors="pt",
    ).to(device)

    outputs = model(**inputs, multimask_output=False)
    masks = processor.post_process_masks(
        outputs.pred_masks.cpu(),
        inputs["original_sizes"],
    )[0]

    return _masks_to_proposals(masks)


def generate_proposals_batch(
    model: Sam2Model,
    processor: Sam2Processor,
    images: Sequence[Image.Image],
    bboxes_per_image: list[list[list[float]]],
) -> list[list[dict[str, Any]]]:
    """Segment a batch of images using SAM2 with bounding box prompts.

    Args:
        model: Sam2Model instance.
        processor: Sam2Processor instance.
        images: Sequence of PIL Images to segment.
        bboxes_per_image: Bounding boxes per image, outer list matches images length.

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

    device = get_device()
    image_rgb_list = [img.convert("RGB") for img in image_list]

    inputs = processor(
        images=image_rgb_list,
        input_boxes=bboxes_per_image,
        return_tensors="pt",
    ).to(device)

    outputs = model(**inputs, multimask_output=False)
    all_masks = processor.post_process_masks(
        outputs.pred_masks.cpu(),
        inputs["original_sizes"],
    )

    return [_masks_to_proposals(image_masks) for image_masks in all_masks]


def detect_objects(
    model: Owlv2ForObjectDetection,
    processor: Owlv2Processor,
    image: Image.Image,
    text_labels: list[str],
    score_threshold: float = 0.25,
    postprocess_threshold: float = 0.1,
) -> list[DetectionResult]:
    """Detect objects in a single image using OWLv2 with text labels.

    Runs OWLv2 object detection on a single image using provided text label
    queries. Applies post-process thresholding first, then filters by
    score_threshold.

    Args:
        model: Owlv2ForObjectDetection instance.
        processor: Owlv2Processor instance.
        image: PIL Image to detect objects in.
        text_labels: List of label groups, e.g., ["cat", "dog"], ["car"].
        score_threshold: Minimum score for final detection results (>=).
        postprocess_threshold: Threshold for processor's post-processing (>=).

    Returns:
        List of DetectionResult dicts with bbox, score, label.
    """
    device = get_device()
    inputs = processor(text=text_labels, images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    target_sizes = [(image.height, image.width)]
    result = processor.post_process_grounded_object_detection(
        outputs=outputs,
        target_sizes=target_sizes,
        threshold=postprocess_threshold,
        text_labels=[text_labels],
    )[0]
    return _to_detection_results(result, score_threshold)


def detect_objects_batch(
    model: Owlv2ForObjectDetection,
    processor: Owlv2Processor,
    images: Sequence[Image.Image],
    text_labels_per_image: list[list[str]],
    score_threshold: float = 0.25,
    postprocess_threshold: float = 0.1,
) -> list[list[DetectionResult]]:
    """Detect objects in a batch of images using OWLv2 with text labels.

    Runs OWLv2 object detection on multiple images using per-image text label
    queries. Processes all images in a single batch inference call.

    Args:
        model: Owlv2ForObjectDetection instance.
        processor: Owlv2Processor instance.
        images: Sequence of PIL Images to detect objects in.
        text_labels_per_image: Text labels per image, outer list matches images length.
            Each inner list is per-image label, e.g., [["cat", "dog"], ["car"]].
        score_threshold: Minimum score for final detection results (>=).
        postprocess_threshold: Threshold for processor's post-processing (>=).

    Returns:
        List of lists of DetectionResult dicts, one inner list per image.

    Raises:
        ValueError: If len(images) != len(text_labels_per_image).
    """
    image_list = list(images)
    if not image_list:
        return []

    if len(image_list) != len(text_labels_per_image):
        raise ValueError(
            f"Length mismatch: {len(image_list)} images, "
            f"{len(text_labels_per_image)} text label groups"
        )

    device = get_device()
    inputs = processor(
        text=text_labels_per_image,
        images=image_list,
        return_tensors="pt",
    ).to(device)
    with torch.no_grad():
        outputs = model(**inputs)

    target_sizes = [(img.height, img.width) for img in image_list]
    results = processor.post_process_grounded_object_detection(
        outputs=outputs,
        target_sizes=target_sizes,
        threshold=postprocess_threshold,
        text_labels=text_labels_per_image,
    )

    return [_to_detection_results(result, score_threshold) for result in results]


def _masks_to_proposals(masks: Any) -> list[dict[str, Any]]:
    """Convert model output masks to list of mask dicts."""
    mask_array = _to_numpy_mask_array(masks)
    if mask_array is None:
        return []

    if mask_array.ndim < 2:
        return []

    if mask_array.ndim == 2:
        mask_array = np.expand_dims(mask_array, axis=0)
    else:
        height, width = mask_array.shape[-2], mask_array.shape[-1]
        mask_array = mask_array.reshape(-1, height, width)

    proposals: list[dict[str, Any]] = []
    for single_mask in mask_array:
        mask_dict = _build_mask_dict(single_mask)
        if mask_dict is not None:
            proposals.append(mask_dict)

    return proposals


def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
    """Convert mask-like object to numpy array."""
    if mask_like is None:
        return None
    if isinstance(mask_like, np.ndarray):
        return mask_like

    import torch

    if isinstance(mask_like, torch.Tensor):
        return mask_like.detach().cpu().numpy()

    return None


def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None:
    """Build a mask dictionary from a 2D boolean numpy array."""
    if mask_array.ndim != 2:
        return None
    segment = mask_array.astype(bool)
    area = int(segment.sum())
    if area <= 0:
        return None

    ys, xs = np.where(segment)
    min_y, max_y = int(ys.min()), int(ys.max())
    min_x, max_x = int(xs.min()), int(xs.max())
    bbox = [min_x, min_y, max_x - min_x + 1, max_y - min_y + 1]

    return {
        "segment": segment,
        "area": area,
        "bbox": bbox,
        "predicted_iou": None,
        "stability_score": None,
    }


def _to_detection_results(
    result: dict[str, Any], score_threshold: float
) -> list[DetectionResult]:
    """Convert OWLv2 post-process result to detection results list.

    Extracts boxes, scores, and text labels from result dict, converts tensors
    to Python native types, and filters by score threshold.

    Args:
        result: OWLv2 post-process result dict with keys 'boxes', 'scores',
            'text_labels'.
        score_threshold: Minimum score to include detection (>= threshold).

    Returns:
        List of DetectionResult dicts with bbox, score, label.
    """
    boxes = result["boxes"]
    scores = result["scores"]
    text_labels = result["text_labels"]

    detections: list[DetectionResult] = []
    for box, score, label in zip(boxes, scores, text_labels):
        score_float = float(score.item())
        if score_float >= score_threshold:
            bbox = [float(v) for v in box.tolist()]
            detections.append(
                {
                    "bbox": bbox,
                    "score": score_float,
                    "label": str(label),
                }
            )

    return detections