"""LeRobot dataset v2.1 format validation helpers."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, Optional

from any4robot.lerobot_editor.constants import (
    DATA_DIR,
    EPISODES_FILE,
    INFO_FILE,
    META_DIR,
    TASKS_FILE,
)
from any4robot.lerobot_editor.file_utils import FileSystemManager


REQUIRED_META_FILES = (INFO_FILE, EPISODES_FILE, TASKS_FILE)


def _read_json(path: Path) -> Dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        return json.load(handle)


def validate_lerobot_v2_1(dataset_path: str | Path, max_episodes: Optional[int] = None) -> Dict[str, Any]:
    """
    Validate a dataset against the LeRobot v2.1 directory layout.

    Args:
        dataset_path: Dataset root directory.
        max_episodes: Optional cap on number of episodes to validate.
    """
    root = Path(dataset_path)
    result: Dict[str, Any] = {
        "valid": True,
        "errors": [],
        "warnings": [],
        "checked_episodes": 0,
        "dataset_path": str(root),
    }

    if not root.exists():
        result["valid"] = False
        result["errors"].append(f"Dataset path does not exist: {root}")
        return result

    meta_path = root / META_DIR
    data_path = root / DATA_DIR

    if not meta_path.exists():
        result["valid"] = False
        result["errors"].append(f"Missing meta directory: {meta_path}")
        return result

    missing_meta = [name for name in REQUIRED_META_FILES if not (meta_path / name).exists()]
    if missing_meta:
        result["valid"] = False
        for name in missing_meta:
            result["errors"].append(f"Missing meta file: {meta_path / name}")
        return result

    if not data_path.exists():
        result["valid"] = False
        result["errors"].append(f"Missing data directory: {data_path}")
        return result

    try:
        info = _read_json(meta_path / INFO_FILE)
    except Exception as exc:
        result["valid"] = False
        result["errors"].append(f"Failed to read info.json: {exc}")
        return result

    if "features" not in info:
        result["warnings"].append("info.json missing 'features' field")

    manager = FileSystemManager(root)
    video_features = info.get("features", {})
    episodes_path = meta_path / EPISODES_FILE

    with episodes_path.open("r", encoding="utf-8") as handle:
        for line in handle:
            if not line.strip():
                continue
            episode = json.loads(line)
            episode_index = episode.get("episode_index")

            if episode_index is None:
                result["valid"] = False
                result["errors"].append("Episode entry missing episode_index")
                break

            paths = manager.get_episode_file_paths(int(episode_index), video_features)
            if not paths["data"].exists():
                result["valid"] = False
                result["errors"].append(f"Missing data file for episode {episode_index}: {paths['data']}")
                break

            result["checked_episodes"] += 1
            if max_episodes and result["checked_episodes"] >= max_episodes:
                break

    return result


__all__ = ["validate_lerobot_v2_1"]
