from pathlib import Path
from typing import Any

from luxonis_ml.data import DatasetIterator

from .base_parser import BaseParser, ParserOutput


class ClassificationDirectoryParser(BaseParser):
    """Parses directory with ClassificationDirectory annotations to LDF.

    Supports two directory structures:

    Split structure with train/valid/test subdirectories::

        dataset_dir/
        ├── train/
        │   ├── class1/
        │   │   ├── img1.jpg
        │   │   ├── img2.jpg
        │   │   └── ...
        │   ├── class2/
        │   └── ...
        ├── valid/
        └── test/

    Flat structure (class subdirectories directly in root,
    random splits applied at parse time)::

        dataset_dir/
        ├── class1/
        │   ├── img1.jpg
        │   └── ...
        ├── class2/
        │   └── ...
        └── info.json  (optional metadata file)

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

    @staticmethod
    def validate_split(split_path: Path) -> dict[str, Any] | None:
        if not split_path.exists():
            return None
        classes = [
            d
            for d in split_path.iterdir()
            if d.is_dir()
            and d.name
            not in {
                "train",
                "valid",
                "test",
                "val",
                "validation",
                "images",
                "labels",
            }
        ]
        if not classes:
            return None
        # For now allow info.json, can be extended to other metadata files
        fnames = [
            f
            for f in split_path.iterdir()
            if f.is_file() and f.name != "info.json"
        ]
        if fnames:
            return None
        return {"class_dir": split_path}

    def from_dir(
        self, dataset_dir: Path
    ) -> tuple[list[Path], list[Path], list[Path]]:
        added_train_imgs = self._parse_split(class_dir=dataset_dir / "train")
        added_val_imgs = self._parse_split(class_dir=dataset_dir / "valid")
        added_test_imgs = self._parse_split(class_dir=dataset_dir / "test")
        return added_train_imgs, added_val_imgs, added_test_imgs

    def from_split(self, class_dir: Path) -> ParserOutput:
        """Parses annotations from classification directory format to.

        LDF. Annotations include classification
        @type class_dir: Path
        @param class_dir: Path to top level directory
        @rtype: L{ParserOutput}
        @return: Annotation generator, list of classes names, skeleton
            dictionary for keypoints and list of added images.
        """
        class_names = [d.name for d in class_dir.iterdir() if d.is_dir()]

        def generator() -> DatasetIterator:
            for class_name in class_names:
                for img_path in self._list_images(class_dir / class_name):
                    yield {
                        "file": str(img_path.absolute().resolve()),
                        "annotation": {"class": class_name},
                    }

        added_images = self._get_added_images(generator())

        return generator(), {}, added_images
