from __future__ import annotations

import os
from pathlib import Path
from xml.etree.ElementTree import Element, SubElement

import cv2
import numpy as np
import numpy.typing as npt
from defusedxml.ElementTree import parse, tostring
from defusedxml.minidom import parseString

from supervision.dataset.utils import approximate_mask_with_polygons
from supervision.detection.core import Detections
from supervision.detection.utils.converters import polygon_to_mask, polygon_to_xyxy
from supervision.utils.file import list_files_with_extensions


def object_to_pascal_voc(
    xyxy: npt.NDArray[np.number],
    name: str,
    polygon: npt.NDArray[np.number] | None = None,
) -> Element:
    """Build a Pascal VOC ``<object>`` XML element for one detection.

    Coordinates are converted to 1-indexed Pascal VOC convention before writing.
    The input arrays are never mutated; new arrays are allocated for the offset.

    Args:
        xyxy: Bounding box in zero-indexed pixel coordinates ``[x1, y1, x2, y2]``.
            Shape ``(4,)``.
        name: Class label string written to the ``<name>`` child element.
        polygon: Optional segmentation polygon in zero-indexed pixel coordinates.
            Shape ``(N, 2)``.

    Returns:
        An XML ``Element`` rooted at ``<object>`` containing ``<name>``,
        ``<bndbox>``, and optionally ``<polygon>`` children.

    Examples:
        >>> import numpy as np
        >>> from supervision.dataset.formats.pascal_voc import object_to_pascal_voc
        >>> elem = object_to_pascal_voc(np.array([0, 0, 9, 9]), name="cat")
        >>> elem.find("bndbox/xmin").text
        '1'
        >>> elem.find("bndbox/xmax").text
        '10'
    """
    root = Element("object")

    object_name = SubElement(root, "name")
    object_name.text = name

    # Pascal VOC coordinates are 1-indexed (https://github.com/roboflow/supervision/issues/144).
    # Rebind to a new array instead of `+= 1`: `xyxy` is a view into the source
    # `Detections.xyxy` (yielded by `Detections.__iter__`), so an in-place add
    # would corrupt the caller's detections by +1 on every export.
    xyxy = xyxy + 1

    bndbox = SubElement(root, "bndbox")
    xmin = SubElement(bndbox, "xmin")
    xmin.text = str(int(xyxy[0]))
    ymin = SubElement(bndbox, "ymin")
    ymin.text = str(int(xyxy[1]))
    xmax = SubElement(bndbox, "xmax")
    xmax.text = str(int(xyxy[2]))
    ymax = SubElement(bndbox, "ymax")
    ymax.text = str(int(xyxy[3]))

    if polygon is not None:
        # 1-indexed, rebound to avoid mutating the caller's array (see above).
        polygon = polygon + 1
        object_polygon = SubElement(root, "polygon")
        for index, point in enumerate(polygon, start=1):
            x_coordinate, y_coordinate = point
            x = SubElement(object_polygon, f"x{index}")
            x.text = str(x_coordinate)
            y = SubElement(object_polygon, f"y{index}")
            y.text = str(y_coordinate)

    return root


def detections_to_pascal_voc(
    detections: Detections,
    classes: list[str],
    filename: str,
    image_shape: tuple[int, int, int],
    min_image_area_percentage: float = 0.0,
    max_image_area_percentage: float = 1.0,
    approximation_percentage: float = 0.75,
) -> str:
    """
    Converts Detections object to Pascal VOC XML format.

    Args:
        detections: A Detections object containing bounding boxes,
            class ids, and other relevant information.
        classes: A list of class names corresponding to the
            class ids in the Detections object.
        filename: The name of the image file associated with the detections.
        image_shape: The shape of the image
            file associated with the detections.
        min_image_area_percentage: Minimum detection area
            relative to area of image associated with it.
        max_image_area_percentage: Maximum detection area
            relative to area of image associated with it.
        approximation_percentage: The percentage of
            polygon points to be removed from the input polygon, in the range [0, 1).
    Returns:
        An XML string in Pascal VOC format representing the detections.

    Note:
        ``detections`` is never mutated by this function; the source ``xyxy``
        array is unchanged after the call. The function is therefore safe to
        call multiple times on the same ``Detections`` object.
    """
    height, width, depth = image_shape

    # Create root element
    annotation = Element("annotation")

    # Add folder element
    folder = SubElement(annotation, "folder")
    folder.text = "VOC"

    # Add filename element
    file_name = SubElement(annotation, "filename")
    file_name.text = filename

    # Add source element
    source = SubElement(annotation, "source")
    database = SubElement(source, "database")
    database.text = "roboflow.ai"

    # Add size element
    size = SubElement(annotation, "size")
    w = SubElement(size, "width")
    w.text = str(width)
    h = SubElement(size, "height")
    h.text = str(height)
    d = SubElement(size, "depth")
    d.text = str(depth)

    # Add segmented element
    segmented = SubElement(annotation, "segmented")
    segmented.text = "0"

    # Add object elements
    for xyxy, mask, _, class_id, _, _ in detections:
        if class_id is None:
            raise ValueError("Detections must include class_id for Pascal VOC export.")
        if not isinstance(class_id, (int, np.integer)):
            raise ValueError(
                f"Detections class_id must be an integer for Pascal VOC export, "
                f"got {type(class_id)!r}."
            )
        name = classes[class_id]
        if mask is not None:
            polygons = approximate_mask_with_polygons(
                mask=mask,
                min_image_area_percentage=min_image_area_percentage,
                max_image_area_percentage=max_image_area_percentage,
                approximation_percentage=approximation_percentage,
            )
            for polygon in polygons:
                xyxy = polygon_to_xyxy(polygon=polygon)
                next_object = object_to_pascal_voc(
                    xyxy=xyxy, name=name, polygon=polygon
                )
                annotation.append(next_object)
        else:
            next_object = object_to_pascal_voc(xyxy=xyxy, name=name)
            annotation.append(next_object)

    # Generate XML string
    xml_string = str(parseString(tostring(annotation)).toprettyxml(indent="  "))
    return xml_string


def load_pascal_voc_annotations(
    images_directory_path: str,
    annotations_directory_path: str,
    force_masks: bool = False,
) -> tuple[list[str], list[str], dict[str, Detections]]:
    """
    Loads PASCAL VOC XML annotations and returns the image name,
        a Detections instance, and a list of class names.

    Args:
        images_directory_path: The path to the directory containing the images.
        annotations_directory_path: The path to the directory containing the
            PASCAL VOC annotation files.
        force_masks: If True, forces masks to be loaded for all
            annotations, regardless of whether they are present.

    Returns:
        A tuple with a list
            of class names, a list of paths to images, and a dictionary with image
            paths as keys and corresponding Detections instances as values.
    """

    image_paths = [
        str(path)
        for path in list_files_with_extensions(
            directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
        )
    ]

    classes: list[str] = []
    annotations = {}

    for image_path in image_paths:
        image_stem = Path(image_path).stem
        annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.xml")
        if not os.path.exists(annotation_path):
            annotations[image_path] = Detections.empty()
            continue

        tree = parse(annotation_path)
        root = tree.getroot()

        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"Could not read image from path: {image_path}")
        resolution_wh = (image.shape[1], image.shape[0])
        annotation, classes = detections_from_xml_obj(
            root, classes, resolution_wh, force_masks
        )
        annotations[image_path] = annotation

    return classes, image_paths, annotations


def detections_from_xml_obj(
    root: Element,
    classes: list[str],
    resolution_wh: tuple[int, int],
    force_masks: bool = False,
) -> tuple[Detections, list[str]]:
    """
    Converts an XML object in Pascal VOC format to a Detections object.
    Expected XML format:
    <annotation>
        ...
        <object>
            <name>dog</name>
            <bndbox>
                <xmin>48</xmin>
                <ymin>240</ymin>
                <xmax>195</xmax>
                <ymax>371</ymax>
            </bndbox>
            <polygon>
                <x1>48</x1>
                <y1>240</y1>
                <x2>195</x2>
                <y2>240</y2>
                <x3>195</x3>
                <y3>371</y3>
                <x4>48</x4>
                <y4>371</y4>
            </polygon>
        </object>
    </annotation>

    Returns:
        A tuple containing a Detections object and an
            updated list of class names, extended with the class names
            from the XML object.
    """
    xyxy: list[list[int]] = []
    class_names: list[str] = []
    masks: list[npt.NDArray[np.bool_]] = []
    with_masks = force_masks or any(
        _with_poly_mask(obj) for obj in root.findall("object")
    )
    extended_classes = classes[:]
    for obj in root.findall("object"):
        class_name = _get_required_text(obj, "name")
        class_names.append(class_name)

        bbox = obj.find("bndbox")
        if bbox is None:
            raise ValueError("Missing bndbox in Pascal VOC annotation.")
        x1 = int(_get_required_text(bbox, "xmin"))
        y1 = int(_get_required_text(bbox, "ymin"))
        x2 = int(_get_required_text(bbox, "xmax"))
        y2 = int(_get_required_text(bbox, "ymax"))

        xyxy.append([x1, y1, x2, y2])

        object_mask: npt.NDArray[np.bool_] = np.zeros(
            (resolution_wh[1], resolution_wh[0]), dtype=bool
        )
        for polygon_element in obj.findall("polygon"):
            polygon = parse_polygon_points(polygon_element)
            # https://github.com/roboflow/supervision/issues/144
            polygon -= 1

            mask_from_polygon = polygon_to_mask(
                polygon=polygon,
                resolution_wh=resolution_wh,
            )
            object_mask |= mask_from_polygon.astype(bool)

        if with_masks:
            masks.append(object_mask)

    xyxy_arr: npt.NDArray[np.float32]
    if xyxy:
        xyxy_arr = np.array(xyxy, dtype=np.float32)
    else:
        xyxy_arr = np.empty((0, 4), dtype=np.float32)

    # https://github.com/roboflow/supervision/issues/144
    xyxy_arr -= 1

    for k in set(class_names):
        if k not in extended_classes:
            extended_classes.append(k)
    class_id = np.array(
        [extended_classes.index(class_name) for class_name in class_names]
    )

    annotation = Detections(
        xyxy=xyxy_arr,
        mask=np.array(masks, dtype=bool) if with_masks else None,
        class_id=class_id,
    )

    return annotation, extended_classes


def _with_poly_mask(obj: Element) -> bool:
    return obj.find("polygon") is not None


def parse_polygon_points(polygon: Element) -> npt.NDArray[np.int_]:
    coordinates: list[int] = []
    for coord in polygon.findall(".//*"):
        if coord.text is None:
            raise ValueError("Missing polygon coordinate value in Pascal VOC.")
        coordinates.append(int(coord.text))
    return np.array(
        [(coordinates[i], coordinates[i + 1]) for i in range(0, len(coordinates), 2)],
        dtype=int,
    )


def _get_required_text(element: Element, tag: str) -> str:
    child = element.find(tag)
    if child is None or child.text is None:
        raise ValueError(f"Missing '{tag}' in Pascal VOC annotation.")
    return child.text