from enum import Enum
from pathlib import Path
from typing import Any

import cv2
import numpy as np
import yaml

from luxonis_ml.data import DatasetIterator

from .base_parser import BaseParser, ParserOutput


class Format(str, Enum):
    ROBOFLOW = "roboflow"
    ULTRALYTICS = "ultralytics"


class YOLOv8Parser(BaseParser):
    """Parses annotations from YOLOv8 format (supports v8-v12 models)
    and Ultralytics format to LDF.

    Expected format::

        dataset_dir/
        ├── images/
        │   ├── train/
        │   │   ├── img1.jpg
        │   │   ├── img2.jpg
        │   │   └── ...
        │   ├── val/
        │   └── test/
        ├── labels/
        │   ├── train/
        │   │   ├── img1.txt
        │   │   ├── img2.txt
        │   │   └── ...
        │   ├── val/
        │   └── test/
        └── *.yaml

        OR::

        dataset_dir/
        ├── train/
        │   ├── images/
        │   │   ├── img1.jpg
        │   │   ├── img2.jpg
        │   │   └── ...
        │   ├── labels/
        │   │   ├── img1.txt
        │   │   ├── img2.txt
        │   │   └── ...
        ├── valid/
        │   ├── images/
        │   │   ├── img1.txt
        │   │   ├── img2.txt
        │   │   └── ...
        │   ├── labels/
        │   │   ├── img1.txt
        │   │   ├── img2.txt
        │   │   └── ...
        ├── test/
        │   ├── images/
        │   │   ├── img1.txt
        │   │   ├── img2.txt
        │   │   └── ...
        │   ├── labels/
        │   │   ├── img1.txt
        │   │   ├── img2.txt
        │   │   └── ...
        └── *.yaml

    C{*.yaml} contains names of all present classes.

    This is one of the formats that can be generated by
    U{Roboflow <https://roboflow.com/>}.
    """

    def fit_boundingbox(self, points: np.ndarray) -> dict[str, float]:
        """Fits a bounding box of the polygon (mask)."""
        x_min = np.min(points[:, 0])
        y_min = np.min(points[:, 1])
        x_max = np.max(points[:, 0])
        y_max = np.max(points[:, 1])
        return {
            "x": x_min,
            "y": y_min,
            "w": x_max - x_min,
            "h": y_max - y_min,
        }

    @staticmethod
    def _detect_dataset_dir_format(
        dataset_dir: Path,
    ) -> tuple[Format | None, list[str]]:
        """Checks if dataset directory structure is in Ultralytics or
        Roboflow format."""
        roboflow_folders = ["train", "valid"]  # test folder is optional
        ultralytics_folders = ["images", "labels"]

        existing = [d.name for d in dataset_dir.iterdir() if d.is_dir()]

        if all(folder in existing for folder in roboflow_folders):
            return Format.ROBOFLOW, existing
        if all(folder in existing for folder in ultralytics_folders):
            return Format.ULTRALYTICS, existing
        return None, []

    @staticmethod
    def validate_split(split_path: Path) -> dict[str, Any] | None:
        if not split_path.exists():
            return None

        candidates = [
            (
                split_path / "images",
                split_path / "labels",
                split_path.parent,
            ),  # ROBOFLOW
            (
                split_path.parent.parent / "images" / split_path.name,
                split_path.parent.parent / "labels" / split_path.name,
                split_path.parent.parent,
            ),  # ULTRALYTICS
        ]

        images_path = labels_path = yaml_root = None
        for img_dir, lbl_dir, yroot in candidates:
            if img_dir.exists() and lbl_dir.exists():
                images_path, labels_path, yaml_root = img_dir, lbl_dir, yroot
                break

        if images_path is None or labels_path is None or yaml_root is None:
            return None

        images = BaseParser._list_images(images_path)
        if not images:
            return None

        yaml_file = next(
            (f for ext in ("*.yaml", "*.yml") for f in yaml_root.glob(ext)),
            None,
        )
        if yaml_file is None:
            return None

        return {
            "image_dir": images_path,
            "annotation_dir": labels_path,
            "classes_path": yaml_file,
        }

    @classmethod
    def validate(cls, dataset_dir: Path) -> bool:
        dir_format, splits = cls._detect_dataset_dir_format(dataset_dir)
        if dir_format is None:
            return False

        yaml_file = next(
            (f for ext in ("*.yaml", "*.yml") for f in dataset_dir.glob(ext)),
            None,
        )
        if not yaml_file:
            return False

        if dir_format is Format.ROBOFLOW:
            splits = [
                d.name
                for d in dataset_dir.iterdir()
                if d.is_dir() and d.name in ("train", "valid", "test")
            ]
            if "train" not in splits or len(splits) < 2:
                return False
            return all(cls.validate_split(dataset_dir / s) for s in splits)

        if dir_format is Format.ULTRALYTICS:
            non_split_folders = ["images", "labels"]
            folders = [d.name for d in dataset_dir.iterdir() if d.is_dir()]
            if not all(f in folders for f in non_split_folders):
                return False
            subfolders = [
                d.name
                for d in (dataset_dir / "images").iterdir()
                if d.is_dir()
            ]
            return all(
                cls.validate_split(dataset_dir / "images" / split)
                for split in subfolders
            )

        return False

    def from_dir(
        self, dataset_dir: Path
    ) -> tuple[list[Path], list[Path], list[Path]]:
        yaml_file = next(
            (f for ext in ("*.yaml", "*.yml") for f in dataset_dir.glob(ext)),
            None,
        )
        if not yaml_file:
            raise ValueError("Exactly one yaml file is expected")
        classes_path = dataset_dir / yaml_file.name
        dir_format, _splits = self._detect_dataset_dir_format(dataset_dir)
        added_train_imgs = self._parse_split(
            image_dir=(
                dataset_dir / "images" / "train"
                if dir_format is Format.ULTRALYTICS
                else dataset_dir / "train" / "images"
            ),
            annotation_dir=(
                dataset_dir / "labels" / "train"
                if dir_format is Format.ULTRALYTICS
                else dataset_dir / "train" / "labels"
            ),
            classes_path=classes_path,
        )
        added_val_imgs = self._parse_split(
            image_dir=(
                dataset_dir / "images" / "val"
                if dir_format is Format.ULTRALYTICS
                else dataset_dir / "valid" / "images"
            ),
            annotation_dir=(
                dataset_dir / "labels" / "val"
                if dir_format is Format.ULTRALYTICS
                else dataset_dir / "valid" / "labels"
            ),
            classes_path=classes_path,
        )
        added_test_imgs = self._parse_split(
            image_dir=(
                dataset_dir / "images" / "test"
                if dir_format is Format.ULTRALYTICS
                else dataset_dir / "test" / "images"
            ),
            annotation_dir=(
                dataset_dir / "labels" / "test"
                if dir_format is Format.ULTRALYTICS
                else dataset_dir / "test" / "labels"
            ),
            classes_path=classes_path,
        )

        return added_train_imgs, added_val_imgs, added_test_imgs

    def from_split(
        self, image_dir: Path, annotation_dir: Path, classes_path: Path
    ) -> ParserOutput:
        """Parses annotations from YoloV8 or Ultralytics format to LDF.
        Annotations include object detection, instance segmentation and
        keypoints.

        @type image_dir: Path
        @param image_dir: Path to directory with images
        @type annotation_dir: Path
        @param annotation_dir: Path to directory with annotations
        @type classes_path: Path
        @param classes_path: Path to yaml file with classes names
        @rtype: L{ParserOutput}
        @return: Annotation generator, list of classes names, skeleton
            dictionary for keypoints and list of added images.
        """
        with open(classes_path) as f:
            classes_data = yaml.safe_load(f)

        if isinstance(classes_data["names"], list):
            """
            names: ["class1", "class2", "class3"]
            """
            class_names = dict(enumerate(classes_data["names"]))
        else:
            """
            names:
                0: class1
                1: class2
                2: class3
            """
            class_names = classes_data["names"]

        def generator() -> DatasetIterator:
            for img_path in self._list_images(image_dir):
                ann_path = annotation_dir / img_path.with_suffix(".txt").name

                annotation_data = []
                # First check if file exists (would crash otherwise)
                if ann_path.exists():
                    with open(ann_path) as f:
                        annotation_data = f.readlines()

                # pre-filter empty lines
                annotation_data = [
                    line for line in annotation_data if line.strip()
                ]

                # Handle missing annotations
                if not annotation_data:
                    yield {"file": str(img_path), "annotation": None}
                    continue

                for instance_id, ann_line in enumerate(annotation_data):
                    annotation_elements = ann_line.split()
                    # object detection format: class_id x_center y_center width height
                    # segmentation format: class_id x1 y1 x2 y2 x3 y3 ... xn yn (min 3 points)
                    # keypoints format: class_id x_center y_center width height kp1_x kp1_y kp2_x kp2_y ... kpn_x kpn_y (it can also have 3rd dimension for visibility)

                    if len(annotation_elements) == 5:
                        task_type = "detection"
                    elif len(annotation_elements) > 5:
                        if classes_data.get("kpt_shape", None) is not None:
                            task_type = "keypoints"
                        else:
                            task_type = "segmentation"

                    if task_type == "detection":
                        class_id, x_center, y_center, width, height = (
                            annotation_elements
                        )
                        class_name = class_names[int(class_id)]

                        yield {
                            "file": str(img_path),
                            "annotation": {
                                "class": class_name,
                                "instance_id": instance_id,
                                "boundingbox": {
                                    "x": float(x_center) - float(width) / 2,
                                    "y": float(y_center) - float(height) / 2,
                                    "w": float(width),
                                    "h": float(height),
                                },
                            },
                        }

                    elif task_type == "segmentation":
                        img = cv2.imread(str(img_path))
                        height, width = img.shape[:2]

                        class_id, *points = annotation_elements
                        points = [float(p) for p in points]
                        points = np.array(points).reshape(-1, 2)
                        boundingbox = self.fit_boundingbox(points)
                        points = [(p[0], p[1]) for p in points.tolist()]
                        class_name = class_names[int(class_id)]

                        yield {
                            "file": str(img_path),
                            "annotation": {
                                "class": class_name,
                                "instance_id": instance_id,
                                "boundingbox": boundingbox,
                                "instance_segmentation": {
                                    "height": height,
                                    "width": width,
                                    "points": points,
                                },
                            },
                        }

                    elif task_type == "keypoints":
                        kpt_shape = classes_data.get("kpt_shape")
                        n_kpts, kpt_dim = kpt_shape
                        class_id, *points = annotation_elements
                        x_center, y_center, width, height, *keypoints = points
                        keypoints = [float(p) for p in keypoints]
                        keypoints = np.array(keypoints).reshape(
                            n_kpts, kpt_dim
                        )
                        class_name = class_names[int(class_id)]

                        if kpt_dim == 2:
                            # add full visibility as last dimension
                            keypoints = np.concatenate(
                                [keypoints, np.ones((n_kpts, 1)) * 2], axis=1
                            )

                        keypoints = [
                            (p[0], p[1], int(p[2])) for p in keypoints.tolist()
                        ]

                        yield {
                            "file": str(img_path),
                            "annotation": {
                                "class": class_name,
                                "instance_id": instance_id,
                                "boundingbox": {
                                    "x": float(x_center) - float(width) / 2,
                                    "y": float(y_center) - float(height) / 2,
                                    "w": float(width),
                                    "h": float(height),
                                },
                                "keypoints": {
                                    "keypoints": keypoints,
                                },
                            },
                        }

                    instance_id += 1

        added_images = self._get_added_images(generator())

        return generator(), {}, added_images
