"""YOLO dataset detection and handling."""

import json
from dataclasses import dataclass, field
from pathlib import Path

import yaml

from argus.core.base import IMAGE_EXTENSIONS, Dataset, DatasetFormat, TaskType


@dataclass
class YOLODataset(Dataset):
    """YOLO format dataset.

    Supports detection, segmentation, and classification tasks.

    Structure (detection/segmentation):
        dataset/
        ├── data.yaml (or *.yaml/*.yml with 'names' key)
        ├── images/
        │   ├── train/
        │   └── val/
        └── labels/
            ├── train/
            └── val/

    Structure (classification):
        dataset/
        ├── images/
        │   ├── train/
        │   │   ├── class1/
        │   │   │   ├── img1.jpg
        │   │   │   └── img2.jpg
        │   │   └── class2/
        │   │       └── img1.jpg
        │   └── val/
        │       ├── class1/
        │       └── class2/
    """

    config_file: Path | None = None
    format: DatasetFormat = field(default=DatasetFormat.YOLO, init=False)
    _roboflow_layout: bool = field(default=False, repr=False)

    @classmethod
    def detect(cls, path: Path) -> "YOLODataset | None":
        """Detect if the given path contains a YOLO dataset.

        Args:
            path: Directory path to check for dataset.

        Returns:
            YOLODataset instance if detected, None otherwise.
        """
        path = Path(path)

        if not path.is_dir():
            return None

        # Try detection/segmentation (YAML-based) first
        result = cls._detect_yaml_based(path)
        if result:
            return result

        # Try classification (directory-based structure)
        return cls._detect_classification(path)

    @classmethod
    def _detect_yaml_based(cls, path: Path) -> "YOLODataset | None":
        """Detect YAML-based YOLO dataset (detection/segmentation).

        Args:
            path: Directory path to check.

        Returns:
            YOLODataset if valid YAML config found, None otherwise.
        """
        # Find all YAML files in the directory
        yaml_files = list(path.glob("*.yaml")) + list(path.glob("*.yml"))

        for yaml_file in yaml_files:
            try:
                with open(yaml_file, encoding="utf-8") as f:
                    config = yaml.safe_load(f)

                if not isinstance(config, dict):
                    continue

                # Must have 'names' key for valid YOLO dataset
                if "names" not in config:
                    continue

                # Skip if this looks like a mask dataset config
                # (has ignore_index or palette keys which are mask-specific)
                if "ignore_index" in config or "palette" in config:
                    continue

                names = config["names"]

                # Extract class names
                if isinstance(names, dict):
                    class_names = list(names.values())
                elif isinstance(names, list):
                    class_names = names
                else:
                    continue

                num_classes = len(class_names)

                # Detect available splits
                splits, is_roboflow = cls._detect_splits(path, config)

                # Determine task type (detection vs segmentation)
                task = cls._determine_task_type(path, config, is_roboflow)

                dataset = cls(
                    path=path,
                    task=task,
                    num_classes=num_classes,
                    class_names=class_names,
                    splits=splits,
                    config_file=yaml_file,
                )
                dataset._roboflow_layout = is_roboflow
                return dataset

            except (yaml.YAMLError, OSError):
                continue

        return None

    @classmethod
    def _detect_classification(cls, path: Path) -> "YOLODataset | None":
        """Detect classification dataset from directory structure.

        Classification datasets can have two structures:

        1. Split structure:
            images/{split}/class_name/image.jpg

        2. Flat structure (unsplit):
            class_name/image.jpg

        No YAML config required - class names inferred from directory names.

        Args:
            path: Directory path to check.

        Returns:
            YOLODataset if classification structure found, None otherwise.
        """

        # Try split structure first: images/{split}/class/
        images_root = path / "images"
        if images_root.is_dir():
            splits: list[str] = []
            class_names_set: set[str] = set()

            for split_name in ["train", "val", "test"]:
                split_dir = images_root / split_name
                if not split_dir.is_dir():
                    continue

                # Get subdirectories (potential class folders)
                class_dirs = [d for d in split_dir.iterdir() if d.is_dir()]
                if not class_dirs:
                    continue

                # Check if at least one class dir contains images
                has_images = False
                for class_dir in class_dirs:
                    for f in class_dir.iterdir():
                        if f.suffix.lower() in IMAGE_EXTENSIONS:
                            has_images = True
                            break
                    if has_images:
                        break

                if has_images:
                    splits.append(split_name)
                    class_names_set.update(d.name for d in class_dirs)

            if splits and class_names_set:
                class_names = sorted(class_names_set)
                return cls(
                    path=path,
                    task=TaskType.CLASSIFICATION,
                    num_classes=len(class_names),
                    class_names=class_names,
                    splits=splits,
                    config_file=None,
                )

        # Try flat structure: class_name/image.jpg (no images/ or split dirs)
        # Check if root contains subdirectories with images
        class_dirs = [d for d in path.iterdir() if d.is_dir()]

        # Filter out common non-class directories
        excluded_dirs = {"images", "labels", "annotations", ".git", "__pycache__"}
        class_dirs = [d for d in class_dirs if d.name not in excluded_dirs]

        if not class_dirs:
            return None

        # Check if any directory contains COCO annotation files
        # (Roboflow COCO datasets have train/valid/test dirs with JSON files)
        for class_dir in class_dirs:
            if cls._has_coco_annotation(class_dir):
                return None  # Not a classification dataset

        # Check if these are class directories (contain images directly)
        class_names_set = set()
        for class_dir in class_dirs:
            has_images = any(
                f.suffix.lower() in IMAGE_EXTENSIONS
                for f in class_dir.iterdir()
                if f.is_file()
            )
            if has_images:
                class_names_set.add(class_dir.name)

        # Need at least 2 classes to be a valid classification dataset
        if len(class_names_set) < 2:
            return None

        class_names = sorted(class_names_set)
        return cls(
            path=path,
            task=TaskType.CLASSIFICATION,
            num_classes=len(class_names),
            class_names=class_names,
            splits=[],  # No splits for flat structure
            config_file=None,
        )

    @staticmethod
    def _has_coco_annotation(directory: Path) -> bool:
        """Check if directory contains COCO annotation files.

        Args:
            directory: Directory to check.

        Returns:
            True if COCO annotation files are found, False otherwise.
        """
        for f in directory.glob("*.json"):
            try:
                with open(f, encoding="utf-8") as fp:
                    data = json.load(fp)
                if isinstance(data, dict) and any(
                    k in data for k in ["images", "annotations", "categories"]
                ):
                    return True
            except (json.JSONDecodeError, OSError):
                pass
        return False

    def get_split_dirs(self, split: str) -> tuple[Path, Path]:
        """Get the images and labels directories for a given split.

        Handles both standard YOLO layout (images/{split}/, labels/{split}/)
        and Roboflow layout ({split}/images/, {split}/labels/).

        Args:
            split: Split name (e.g. "train", "val", "test").

        Returns:
            Tuple of (images_dir, labels_dir).
        """
        if self._roboflow_layout:
            # Roboflow uses "valid" instead of "val"
            actual_dir = split
            if split == "val" and not (self.path / "val").is_dir():
                actual_dir = "valid"
            return (
                self.path / actual_dir / "images",
                self.path / actual_dir / "labels",
            )
        # Standard layout — also check "valid" alias for "val"
        actual_split = split
        if (
            split == "val"
            and not (self.path / "images" / "val").is_dir()
            and (self.path / "images" / "valid").is_dir()
        ):
            actual_split = "valid"
        return (
            self.path / "images" / actual_split,
            self.path / "labels" / actual_split,
        )

    def get_instance_counts(self) -> dict[str, dict[str, int]]:
        """Get the number of annotation instances per class, per split.

        For detection/segmentation: Scans images and reads the corresponding
        label file (derived by replacing ``images`` with ``labels`` and the
        extension with ``.txt``).  Orphaned label files without a matching
        image are ignored, matching ultralytics behaviour.

        For classification: Counts images in each class directory
        (1 image = 1 instance).

        Returns:
            Dictionary mapping split name to dict of class name to instance count.
        """
        # Handle classification datasets differently
        if self.task == TaskType.CLASSIFICATION:
            return self._get_classification_instance_counts()

        counts: dict[str, dict[str, int]] = {}

        # Build class_id -> class_name mapping
        id_to_name = {i: name for i, name in enumerate(self.class_names)}

        # Determine which image dirs actually exist
        has_split_image_dirs = (
            any(self.get_split_dirs(s)[0].is_dir() for s in self.splits)
            if self.splits
            else False
        )

        # If splits are declared but no image folders exist, treat as unsplit.
        if self.splits and not has_split_image_dirs:
            splits_to_process = ["unsplit"]
        else:
            splits_to_process = self.splits if self.splits else ["unsplit"]

        for split in splits_to_process:
            split_counts: dict[str, int] = {}

            if split == "unsplit":
                image_dir = self.path / "images"
                label_dir = self.path / "labels"
            else:
                image_dir, label_dir = self.get_split_dirs(split)

            if not image_dir.is_dir():
                continue

            # Iterate over images and read corresponding label files
            for img_file in image_dir.iterdir():
                if img_file.suffix.lower() not in IMAGE_EXTENSIONS:
                    continue

                label_path = label_dir / (img_file.stem + ".txt")
                if not label_path.is_file():
                    continue

                try:
                    with open(label_path, encoding="utf-8") as f:
                        for line in f:
                            line = line.strip()
                            if not line:
                                continue
                            parts = line.split()
                            if len(parts) >= 5:  # Valid annotation line
                                try:
                                    class_id = int(parts[0])
                                    fallback = f"class_{class_id}"
                                    class_name = id_to_name.get(class_id, fallback)
                                    current = split_counts.get(class_name, 0)
                                    split_counts[class_name] = current + 1
                                except ValueError:
                                    continue
                except OSError:
                    continue

            counts[split] = split_counts

        return counts

    def _get_classification_instance_counts(self) -> dict[str, dict[str, int]]:
        """Get instance counts for classification datasets.

        Each image is one instance of its class.

        Returns:
            Dictionary mapping split name to dict of class name to image count.
        """
        counts: dict[str, dict[str, int]] = {}

        # Handle flat structure (no splits)
        if not self.splits:
            split_counts: dict[str, int] = {}
            for class_name in self.class_names:
                class_dir = self.path / class_name
                if not class_dir.is_dir():
                    split_counts[class_name] = 0
                    continue

                image_count = sum(
                    1
                    for f in class_dir.iterdir()
                    if f.suffix.lower() in IMAGE_EXTENSIONS
                )
                split_counts[class_name] = image_count

            counts["unsplit"] = split_counts
            return counts

        # Handle split structure
        images_root = self.path / "images"
        for split in self.splits:
            split_dir = images_root / split
            if not split_dir.is_dir():
                continue

            split_counts = {}
            for class_name in self.class_names:
                class_dir = split_dir / class_name
                if not class_dir.is_dir():
                    split_counts[class_name] = 0
                    continue

                image_count = sum(
                    1
                    for f in class_dir.iterdir()
                    if f.suffix.lower() in IMAGE_EXTENSIONS
                )
                split_counts[class_name] = image_count

            counts[split] = split_counts

        return counts

    def get_image_counts(self) -> dict[str, dict[str, int]]:
        """Get image counts per split, including background images.

        For detection/segmentation: Scans the images directory and derives
        the label path for each image.  Images whose label file is missing
        or empty are counted as background, matching ultralytics behaviour.

        For classification: Counts total images across all class directories.
        Background count is always 0 (no background concept in classification).

        Returns:
            Dictionary mapping split name to dict with "total" and "background" counts.
        """
        # Handle classification datasets differently
        if self.task == TaskType.CLASSIFICATION:
            return self._get_classification_image_counts()

        counts: dict[str, dict[str, int]] = {}

        has_split_image_dirs = (
            any(self.get_split_dirs(s)[0].is_dir() for s in self.splits)
            if self.splits
            else False
        )

        # If splits are declared but no images/{split} folders exist, treat as unsplit.
        if self.splits and not has_split_image_dirs:
            splits_to_process = ["unsplit"]
        else:
            splits_to_process = self.splits if self.splits else ["unsplit"]

        for split in splits_to_process:
            if split == "unsplit":
                image_dir = self.path / "images"
                label_dir = self.path / "labels"
            else:
                image_dir, label_dir = self.get_split_dirs(split)

            if not image_dir.is_dir():
                continue

            total = 0
            background = 0
            for img_file in image_dir.iterdir():
                if img_file.suffix.lower() not in IMAGE_EXTENSIONS:
                    continue

                total += 1
                label_path = label_dir / (img_file.stem + ".txt")
                if not label_path.is_file():
                    background += 1
                    continue
                try:
                    content = label_path.read_text(encoding="utf-8").strip()
                    if not content:
                        background += 1
                except OSError:
                    background += 1

            counts[split] = {"total": total, "background": background}

        return counts

    def _get_classification_image_counts(self) -> dict[str, dict[str, int]]:
        """Get image counts for classification datasets.

        Returns:
            Dictionary mapping split name to dict with "total" and "background" counts.
            Background is always 0 for classification.
        """
        counts: dict[str, dict[str, int]] = {}

        # Handle flat structure (no splits)
        if not self.splits:
            total = 0
            for class_name in self.class_names:
                class_dir = self.path / class_name
                if not class_dir.is_dir():
                    continue

                total += sum(
                    1
                    for f in class_dir.iterdir()
                    if f.suffix.lower() in IMAGE_EXTENSIONS
                )

            counts["unsplit"] = {"total": total, "background": 0}
            return counts

        # Handle split structure
        images_root = self.path / "images"
        for split in self.splits:
            split_dir = images_root / split
            if not split_dir.is_dir():
                continue

            total = 0
            for class_name in self.class_names:
                class_dir = split_dir / class_name
                if not class_dir.is_dir():
                    continue

                total += sum(
                    1
                    for f in class_dir.iterdir()
                    if f.suffix.lower() in IMAGE_EXTENSIONS
                )

            counts[split] = {"total": total, "background": 0}

        return counts

    @classmethod
    def _detect_splits(cls, path: Path, config: dict) -> tuple[list[str], bool]:
        """Detect available splits from config and filesystem.

        Args:
            path: Dataset root path.
            config: Parsed YAML config.

        Returns:
            Tuple of (list of available split names, is_roboflow_layout).
        """
        splits: list[str] = []
        is_roboflow = False

        # Detect Roboflow layout: {split}/images/ exists at the top level
        roboflow_dirs = {"train", "val", "valid", "test"}
        for d in roboflow_dirs:
            if (path / d / "images").is_dir():
                is_roboflow = True
                break

        # Config key aliases: Roboflow uses "val" key pointing to "../valid/images"
        config_key_aliases = {
            "train": ["train"],
            "val": ["val", "valid"],
            "test": ["test"],
        }

        for split_name, aliases in config_key_aliases.items():
            found = False

            # Check config-defined paths
            for alias in aliases:
                if alias in config:
                    config_val = config[alias]
                    if config_val:
                        full_path = path / config_val
                        if full_path.exists():
                            splits.append(split_name)
                            found = True
                            break

            if found:
                continue

            if is_roboflow:
                # Roboflow layout: {split}/images/
                # Check "valid" as alias for "val"
                dirs_to_check = (
                    [split_name, "valid"] if split_name == "val" else [split_name]
                )
                for d in dirs_to_check:
                    if (path / d / "images").is_dir():
                        splits.append(split_name)
                        found = True
                        break
                if found:
                    continue
            else:
                # Standard layout: images/{split}/ (also check "valid" alias)
                dirs_to_check = (
                    [split_name, "valid"] if split_name == "val" else [split_name]
                )
                for alias in dirs_to_check:
                    if (path / "images" / alias).is_dir():
                        splits.append(split_name)
                        found = True
                        break
                if found:
                    continue

                # Pattern: train/, val/ (flat structure without images/ or labels/)
                if (path / split_name).is_dir():
                    if (path / "images").is_dir() or (path / "labels").is_dir():
                        continue
                    splits.append(split_name)

        return splits, is_roboflow

    @classmethod
    def _determine_task_type(
        cls, path: Path, config: dict, is_roboflow: bool = False
    ) -> TaskType:
        """Determine if dataset is detection or segmentation.

        Detection labels have 5 columns: class x_center y_center width height
        Segmentation labels have >5 columns: class x1 y1 x2 y2 ... xn yn

        Args:
            path: Dataset root path.
            config: Parsed YAML config.
            is_roboflow: Whether the dataset uses Roboflow layout.

        Returns:
            TaskType.DETECTION or TaskType.SEGMENTATION.
        """
        # Find label files to analyze
        label_dirs = []

        # Check common label locations
        for split in ["train", "val", "valid", "test"]:
            # Standard layout: labels/train/
            label_dir = path / "labels" / split
            if label_dir.is_dir():
                label_dirs.append(label_dir)

            # Roboflow layout: train/labels/
            if is_roboflow:
                label_dir = path / split / "labels"
                if label_dir.is_dir():
                    label_dirs.append(label_dir)

        # Pattern: labels/ (flat)
        labels_dir = path / "labels"
        if labels_dir.is_dir() and not label_dirs:
            label_dirs.append(labels_dir)

        # Sample label files
        for label_dir in label_dirs:
            txt_files = list(label_dir.glob("*.txt"))
            for txt_file in txt_files[:5]:  # Sample up to 5 files
                try:
                    with open(txt_file, encoding="utf-8") as f:
                        for line in f:
                            line = line.strip()
                            if not line:
                                continue
                            parts = line.split()
                            if len(parts) > 5:
                                return TaskType.SEGMENTATION
                            elif len(parts) == 5:
                                return TaskType.DETECTION
                except OSError:
                    continue

        # Default to detection if no labels found or inconclusive
        return TaskType.DETECTION

    def get_image_paths(self, split: str | None = None) -> list[Path]:
        """Get all image file paths for a split or the entire dataset.

        Args:
            split: Specific split to get images from. If None, returns all images.

        Returns:
            List of image file paths sorted alphabetically.
        """
        # Handle classification datasets differently
        if self.task == TaskType.CLASSIFICATION:
            return self._get_classification_image_paths(split)

        images_root = self.path / "images"
        image_paths: list[Path] = []
        seen: set[Path] = set()

        # Decide how to interpret splits. Check if split image dirs exist
        # using get_split_dirs for layout-aware resolution.
        has_split_image_dirs = (
            any(self.get_split_dirs(s)[0].is_dir() for s in self.splits)
            if self.splits
            else False
        )

        if split:
            splits_to_search = [split]
        elif self.splits and not has_split_image_dirs:
            splits_to_search = ["unsplit"]
        elif self.splits:
            splits_to_search = self.splits
        else:
            splits_to_search = ["unsplit"]

        for s in splits_to_search:
            if s == "unsplit":
                image_dir = images_root
            else:
                image_dir = self.get_split_dirs(s)[0]
                if not image_dir.is_dir():
                    # If a split was explicitly requested but the folder doesn't
                    # exist, fall back to images/.
                    image_dir = images_root

            if not image_dir.is_dir():
                continue

            for img_file in image_dir.iterdir():
                if img_file.suffix.lower() not in IMAGE_EXTENSIONS:
                    continue

                resolved = img_file.resolve()
                if resolved in seen:
                    continue
                seen.add(resolved)
                image_paths.append(img_file)

        return sorted(image_paths, key=lambda p: p.name)

    def _get_classification_image_paths(self, split: str | None = None) -> list[Path]:
        """Get image paths for classification datasets.

        Args:
            split: Specific split to get images from. If None, returns all images.

        Returns:
            List of image file paths sorted alphabetically.
        """

        image_paths: list[Path] = []

        # Handle flat structure (no splits)
        if not self.splits:
            for class_name in self.class_names:
                class_dir = self.path / class_name
                if not class_dir.is_dir():
                    continue

                for img_file in class_dir.iterdir():
                    if img_file.suffix.lower() in IMAGE_EXTENSIONS:
                        image_paths.append(img_file)

            return sorted(image_paths, key=lambda p: p.name)

        # Handle split structure
        images_root = self.path / "images"
        splits_to_search = [split] if split else self.splits

        for s in splits_to_search:
            split_dir = images_root / s
            if not split_dir.is_dir():
                continue

            for class_name in self.class_names:
                class_dir = split_dir / class_name
                if not class_dir.is_dir():
                    continue

                for img_file in class_dir.iterdir():
                    if img_file.suffix.lower() in IMAGE_EXTENSIONS:
                        image_paths.append(img_file)

        return sorted(image_paths, key=lambda p: p.name)

    def get_annotations_for_image(self, image_path: Path) -> list[dict]:
        """Get annotations for a specific image.

        Args:
            image_path: Path to the image file.

        Returns:
            List of annotation dicts with bbox/polygon in absolute coordinates.
        """
        import cv2

        annotations: list[dict] = []

        # Build class_id -> class_name mapping
        id_to_name = {i: name for i, name in enumerate(self.class_names)}

        # Find the corresponding label file
        # Image: images/train/img.jpg -> Label: labels/train/img.txt
        image_parts = image_path.parts
        try:
            images_idx = image_parts.index("images")
            label_parts = list(image_parts)
            label_parts[images_idx] = "labels"
            label_path = Path(*label_parts).with_suffix(".txt")
        except ValueError:
            # Fallback: look in labels directory with same name
            label_path = self.path / "labels" / image_path.with_suffix(".txt").name

        if not label_path.exists():
            return annotations

        # Get image dimensions for converting normalized coords to absolute
        img = cv2.imread(str(image_path))
        if img is None:
            return annotations
        img_height, img_width = img.shape[:2]

        try:
            with open(label_path, encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue

                    parts = line.split()
                    if len(parts) < 5:
                        continue

                    try:
                        class_id = int(parts[0])
                        class_name = id_to_name.get(class_id, f"class_{class_id}")

                        if len(parts) == 5:
                            # Detection: class x_center y_center width height
                            x_center = float(parts[1]) * img_width
                            y_center = float(parts[2]) * img_height
                            width = float(parts[3]) * img_width
                            height = float(parts[4]) * img_height

                            # Convert to x, y, w, h (top-left corner)
                            x = x_center - width / 2
                            y = y_center - height / 2

                            annotations.append(
                                {
                                    "class_name": class_name,
                                    "class_id": class_id,
                                    "bbox": (x, y, width, height),
                                    "polygon": None,
                                }
                            )
                        else:
                            # Segmentation: class x1 y1 x2 y2 ... xn yn
                            coords = [float(p) for p in parts[1:]]
                            polygon = []
                            for i in range(0, len(coords), 2):
                                px = coords[i] * img_width
                                py = coords[i + 1] * img_height
                                polygon.append((px, py))

                            # Calculate bounding box from polygon
                            xs = [p[0] for p in polygon]
                            ys = [p[1] for p in polygon]
                            x = min(xs)
                            y = min(ys)
                            width = max(xs) - x
                            height = max(ys) - y

                            annotations.append(
                                {
                                    "class_name": class_name,
                                    "class_id": class_id,
                                    "bbox": (x, y, width, height),
                                    "polygon": polygon,
                                }
                            )

                    except (ValueError, IndexError):
                        continue

        except OSError:
            pass

        return annotations

    def get_images_by_class(self, split: str | None = None) -> dict[str, list[Path]]:
        """Get images grouped by class for classification datasets.

        Args:
            split: Specific split to get images from. If None, uses first
                   available split or all images for flat structure.

        Returns:
            Dictionary mapping class name to list of image paths.
        """
        if self.task != TaskType.CLASSIFICATION:
            return {}

        images_by_class: dict[str, list[Path]] = {cls: [] for cls in self.class_names}

        # Handle flat structure (no splits)
        if not self.splits:
            for class_name in self.class_names:
                class_dir = self.path / class_name
                if not class_dir.is_dir():
                    continue

                for img_file in class_dir.iterdir():
                    if img_file.suffix.lower() in IMAGE_EXTENSIONS:
                        images_by_class[class_name].append(img_file)

            # Sort images within each class
            for class_name in images_by_class:
                images_by_class[class_name] = sorted(
                    images_by_class[class_name], key=lambda p: p.name
                )

            return images_by_class

        # Handle split structure
        images_root = self.path / "images"
        default_splits = self.splits[:1] if self.splits else []
        splits_to_search = [split] if split else default_splits

        for s in splits_to_search:
            split_dir = images_root / s
            if not split_dir.is_dir():
                continue

            for class_name in self.class_names:
                class_dir = split_dir / class_name
                if not class_dir.is_dir():
                    continue

                for img_file in class_dir.iterdir():
                    if img_file.suffix.lower() in IMAGE_EXTENSIONS:
                        images_by_class[class_name].append(img_file)

        # Sort images within each class for consistent ordering
        for class_name in images_by_class:
            images_by_class[class_name] = sorted(
                images_by_class[class_name], key=lambda p: p.name
            )

        return images_by_class
