from __future__ import annotations

import io
from typing import Any, cast

import numpy as np
import numpy.typing as npt
from PIL import Image

from supervision.config import CLASS_NAME_DATA_FIELD
from supervision.detection.utils.converters import mask_to_xyxy


def process_transformers_detection_result(
    detection_result: dict[str, Any], id2label: dict[int, str] | None
) -> dict[str, Any]:
    """
    Process the result of Transformers object detection functions such as
    `post_process` (v4) and `post_process_detection` (v5).

    Args:
        detection_result: Dictionary containing detection results with keys
            'boxes', 'labels', and 'scores'.
        id2label: A dictionary mapping class IDs to labels,
            typically part of the `transformers` model configuration. If provided, the
            resulting dictionary will include class names.

    Returns:
        Processed detection result including bounding boxes, confidence scores,
            class IDs, and data.
    """
    class_ids = detection_result["labels"].cpu().detach().numpy().astype(int)
    data = append_class_names_to_data(class_ids, id2label, {})

    return dict(
        xyxy=detection_result["boxes"].cpu().detach().numpy(),
        confidence=detection_result["scores"].cpu().detach().numpy(),
        class_id=class_ids,
        data=data,
    )


def process_transformers_v4_segmentation_result(
    segmentation_result: dict[str, Any], id2label: dict[int, str] | None
) -> dict[str, Any]:
    """
    Process the result of Transformers segmentation functions such as
    `post_process_panoptic`, `post_process_segmentation`, and `post_process_instance`
    (v4).

    Args:
        segmentation_result: Dictionary containing segmentation results with keys
            'masks', 'labels', and 'scores'.
        id2label: A dictionary mapping class IDs to labels,
            typically part of the `transformers` model configuration. If provided, the
            resulting dictionary will include class names.

    Returns:
        Processed segmentation result including bounding boxes, masks, confidence
            scores, class IDs, and data.
    """
    if "png_string" in segmentation_result:
        return process_transformers_v4_panoptic_segmentation_result(
            segmentation_result, id2label
        )
    else:
        boxes = None
        if "boxes" in segmentation_result:
            boxes = segmentation_result["boxes"].cpu().detach().numpy()
        masks = segmentation_result["masks"].cpu().detach().numpy().astype(bool)
        class_ids = segmentation_result["labels"].cpu().detach().numpy().astype(int)

        return dict(
            xyxy=boxes if boxes is not None else mask_to_xyxy(masks),
            mask=np.squeeze(masks, axis=1) if boxes is not None else masks,
            confidence=segmentation_result["scores"].cpu().detach().numpy(),
            class_id=class_ids,
            data=append_class_names_to_data(class_ids, id2label, {}),
        )


def process_transformers_v5_segmentation_result(
    segmentation_result: Any, id2label: dict[int, str] | None
) -> dict[str, Any]:
    """
    Process the result of Transformers segmentation functions such as
    `post_process_semantic_segmentation`, `post_process_instance_segmentation`, and
    `post_process_panoptic_segmentation` (v5).

    Args:
        segmentation_result: Either a dictionary containing segmentation results
            (`segments_info` and `segmentation`) or a tensor object
            representing a panoptic segmentation map.
        id2label: A dictionary mapping class IDs to labels,
            typically part of the `transformers` model configuration. If provided, the
            resulting dictionary will include class names.

    Returns:
        Processed segmentation result including bounding boxes, masks, confidence
            scores, class IDs, and data.
    """
    if segmentation_result.__class__.__name__ == "Tensor":
        segmentation_array = segmentation_result.cpu().detach().numpy()
        return process_transformers_v5_panoptic_segmentation_result(
            segmentation_array, id2label
        )
    return process_transformers_v5_semantic_or_instance_segmentation_result(
        cast(dict[str, Any], segmentation_result), id2label
    )


def process_transformers_v5_semantic_or_instance_segmentation_result(
    segmentation_result: dict[str, Any], id2label: dict[int, str] | None
) -> dict[str, Any]:
    """
    Process the result of Transformers segmentation functions such as
    `post_process_semantic_segmentation` and `post_process_instance_segmentation` (v5).

    Args:
        segmentation_result: Dictionary containing segmentation results with keys
            `segments_info` and `segmentation`.
        id2label: A dictionary mapping class IDs to labels,
            typically part of the `transformers` model configuration. If provided, the
            resulting dictionary will include class names.

    Returns:
        Processed segmentation result including bounding boxes, masks, confidence
            scores, class IDs, and data.
    """
    segments_info = segmentation_result["segments_info"]
    scores = np.array([segment["score"] for segment in segments_info])
    class_ids = np.array([segment["label_id"] for segment in segments_info])
    segmentation_array = segmentation_result["segmentation"].cpu().detach().numpy()
    masks = np.array(
        [segmentation_array == segment["id"] for segment in segments_info]
    ).astype(bool)
    data = append_class_names_to_data(class_ids, id2label, {})

    return dict(
        xyxy=mask_to_xyxy(masks),
        mask=masks,
        confidence=scores,
        class_id=class_ids,
        data=data,
    )


def process_transformers_v4_panoptic_segmentation_result(
    segmentation_result: dict[str, Any], id2label: dict[int, str] | None
) -> dict[str, Any]:
    """
    Process the result of the Transformers function `post_process_panoptic` (v4).

    Args:
        segmentation_result: Dictionary containing segmentation results with keys
            such as 'png_string' and 'segments_info'.
        id2label: A dictionary mapping class IDs to labels,
            typically part of the `transformers` model configuration. If provided, the
            resulting dictionary will include class names.

    Returns:
        Processed segmentation result including bounding boxes, masks,
            class IDs, and data.
    """
    segments_info = segmentation_result["segments_info"]
    png_string = segmentation_result["png_string"]
    class_ids = np.array([segment["category_id"] for segment in segments_info])
    segmentation_array = png_string_to_segmentation_array(png_string=png_string)
    masks = np.array(
        [segmentation_array == segment["id"] for segment in segments_info]
    ).astype(bool)
    data = append_class_names_to_data(class_ids, id2label, {})

    return dict(
        xyxy=mask_to_xyxy(masks),
        mask=masks,
        class_id=class_ids,
        data=data,
    )


def process_transformers_v5_panoptic_segmentation_result(
    segmentation_array: npt.NDArray[Any], id2label: dict[int, str] | None
) -> dict[str, Any]:
    """
    Process the result of the Transformers function
    `post_process_panoptic_segmentation` (v5).

    Args:
        segmentation_array: Segmentation array.
        id2label: A dictionary mapping class IDs to labels,
            typically part of the `transformers` model configuration. If provided, the
            resulting dictionary will include class names.

    Returns:
        Processed segmentation result including bounding boxes, masks,
            class IDs, and data.
    """
    class_ids = np.unique(segmentation_array)
    masks = np.stack(
        [segmentation_array == class_id for class_id in class_ids], axis=0
    ).astype(bool)
    data = append_class_names_to_data(class_ids, id2label, {})
    return dict(xyxy=mask_to_xyxy(masks), mask=masks, class_id=class_ids, data=data)


def png_string_to_segmentation_array(png_string: bytes) -> npt.NDArray[Any]:
    """
    Convert a PNG byte string to a label mask array.

    Args:
        png_string: A byte string representing the PNG image.

    Returns:
        A label mask array with shape (H, W), where H and W
            are the height and width of the image. Each unique value in the array
            represents a different object or category.
    """
    image = Image.open(io.BytesIO(png_string))
    mask = np.array(image, dtype=np.uint8)
    return cast(npt.NDArray[Any], mask[:, :, 0])


def append_class_names_to_data(
    class_ids: npt.NDArray[Any],
    id2label: dict[int, str] | None,
    data: dict[str, Any] | None = None,
) -> dict[str, Any]:
    """
    Helper function to create or append to a data dictionary with class names if
    available.

    Args:
        class_ids: Array of class IDs.
        id2label: A dictionary mapping class IDs to labels,
            typically part of the `transformers` model configuration. If provided, the
            resulting dictionary will include class names.
        data: An existing data dictionary to append to.

    Returns:
        Dictionary containing class names if id2label is provided.
    """
    if data is None:
        data = {}

    if id2label is not None:
        class_names = np.array([id2label[class_id] for class_id in class_ids])
        data[CLASS_NAME_DATA_FIELD] = class_names

    return data