"""Dataset combination utilities."""

from __future__ import annotations

from pathlib import Path
from typing import Dict, Iterable, List, Optional

from any4robot.lerobot_editor.operations_v21 import DatasetOperationsV21
from any4robot.lerobot_editor.constants import EPISODES_FILE, INFO_FILE, META_DIR, TASKS_FILE


def _is_lerobot_dataset(path: Path) -> bool:
    """Return True if path looks like a LeRobot dataset root."""
    meta_dir = path / META_DIR
    return (
        meta_dir.is_dir()
        and (meta_dir / INFO_FILE).exists()
        and (meta_dir / EPISODES_FILE).exists()
        and (meta_dir / TASKS_FILE).exists()
    )


def discover_lerobot_datasets(source_root: str | Path) -> List[Path]:
    """
    Discover LeRobot dataset roots under a folder.

    If `source_root` is itself a dataset, it returns `[source_root]`.
    Otherwise, it returns direct child folders that look like datasets.
    """
    root = Path(source_root)
    if not root.exists():
        raise FileNotFoundError(f"source_root does not exist: {root}")
    if not root.is_dir():
        raise NotADirectoryError(f"source_root is not a directory: {root}")

    if _is_lerobot_dataset(root):
        return [root]

    datasets = [path for path in root.iterdir() if path.is_dir() and _is_lerobot_dataset(path)]
    return sorted(datasets)


def merge_lerobot_datasets(
    source_datasets: Iterable[str | Path],
    output_path: str | Path,
    task_mapping: Optional[Dict[str, str]] = None,
    dry_run: bool = False,
) -> bool:
    """
    Merge multiple LeRobot datasets into a single dataset.

    This delegates to the LERO DatasetOperations merge implementation.
    """
    source_paths = [Path(p) for p in source_datasets]
    output = Path(output_path)

    if not source_paths:
        raise ValueError("source_datasets must contain at least one dataset path")

    # Use the first dataset path to bootstrap operations for merging.
    ops = DatasetOperationsV21(source_paths[0])
    return ops.merge_datasets(source_paths, output, task_mapping=task_mapping, dry_run=dry_run)


def merge_lerobot_datasets_from_folder(
    source_root: str | Path,
    output_path: str | Path,
    task_mapping: Optional[Dict[str, str]] = None,
    dry_run: bool = False,
) -> bool:
    """
    Discover and merge all LeRobot datasets under a folder.
    """
    source_paths = discover_lerobot_datasets(source_root)
    output = Path(output_path)
    output_resolved = output.resolve()

    source_paths = [path for path in source_paths if path.resolve() != output_resolved]
    if not source_paths:
        raise ValueError(
            f"No LeRobot datasets found under source_root: {source_root}"
        )

    return merge_lerobot_datasets(
        source_datasets=source_paths,
        output_path=output,
        task_mapping=task_mapping,
        dry_run=dry_run,
    )


__all__ = [
    "discover_lerobot_datasets",
    "merge_lerobot_datasets",
    "merge_lerobot_datasets_from_folder",
]
