"""Helpers for detection tests."""
import os
import xml.etree.cElementTree as ET  # nosec
from glob import glob
from typing import List, Tuple

import cv2
import numpy as np


class BBFromMasks:
    """Creates temporary XML files from masks for testing. Intended to be used
    as a context so that the XML files are automatically deleted when the
    execution goes out of scope.

    Example:

        >>> with BBFromMasks(root="/tmp/datasets/MVTec", datast_name="MVTec"):
        >>>     tests_case()

    Args:
        root (str, optional): Path to the dataset location. Defaults to "datasets/MVTec".
        dataset_name (str, optional): Name of the dataset to write to the XML file. Defaults to "MVTec".
    """

    def __init__(self, root: str = "datasets/MVTec", dataset_name: str = "MVTec") -> None:
        self.root = root
        self.dataset_name = dataset_name
        self.generated_xml_files: List[str] = []

    def __enter__(self):
        """Generate XML files."""
        for mask_path in glob(os.path.join(self.root, "*/ground_truth/*/*_mask.png")):
            path_tree = mask_path.split("/")
            image = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
            image = np.array(image, dtype=np.uint8)
            im_size = image.shape
            contours, _ = cv2.findContours(image, 1, 1)

            boxes = []
            for contour in contours:
                p1 = [np.min(contour[..., 0]), np.min(contour[..., 1])]
                p2 = [np.max(contour[..., 0]), np.max(contour[..., 1])]
                boxes.append([p1, p2])

            contents = self._create_xml_contents(boxes, path_tree, im_size)
            tree = ET.ElementTree(contents)
            output_loc = "/".join(path_tree[:-1]) + f"/{path_tree[-1].rstrip('_mask.png')}.xml"
            tree.write(output_loc)
            # write the xml
            self.generated_xml_files.append(output_loc)

    def __exit__(self, _exc_type, _exc_value, _exc_traceback):
        """Cleans up generated XML files."""
        for file in self.generated_xml_files:
            os.remove(file)

    def _create_xml_contents(
        self, boxes: List[List[List[np.int]]], path_tree: List[str], image_size: Tuple[int, int]
    ) -> ET.Element:
        """Create the contents of the annotation file in Pascal VOC format.

        Args:
            boxes (List[List[List[np.int]]]): The calculated pox corners from the masks
            path_tree (List[str]): The entire filepath of the mask.png image split into a list
            image_size (Tuple[int, int]): Tuple of image size for writing into annotation

        Returns:
            ET.Element: annotation root element
        """
        annotation = ET.Element("annotation")
        ET.SubElement(annotation, "folder").text = path_tree[-2]
        ET.SubElement(annotation, "filename").text = path_tree[-1]

        source = ET.SubElement(annotation, "source")
        ET.SubElement(source, "database").text = self.dataset_name
        ET.SubElement(source, "annotation").text = "PASCAL VOC"

        size = ET.SubElement(annotation, "size")
        ET.SubElement(size, "width").text = str(image_size[0])
        ET.SubElement(size, "height").text = str(image_size[1])
        ET.SubElement(size, "depth").text = "1"
        for box in boxes:
            object = ET.SubElement(annotation, "object")
            ET.SubElement(object, "name").text = "anomaly"
            ET.SubElement(object, "difficult").text = "1"
            bndbox = ET.SubElement(object, "bndbox")
            ET.SubElement(bndbox, "xmin").text = str(box[0][0])
            ET.SubElement(bndbox, "ymin").text = str(box[0][1])
            ET.SubElement(bndbox, "xmax").text = str(box[1][0])
            ET.SubElement(bndbox, "ymax").text = str(box[1][1])

        return annotation