"""Episode-level quality checks and cleanup helpers."""

from __future__ import annotations

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

from any4robot.lerobot_editor.constants import EPISODES_FILE, INFO_FILE, META_DIR
from any4robot.lerobot_editor.core import LeRobotDatasetEditor
from any4robot.lerobot_editor.file_utils import FileSystemManager


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


def _iter_episode_rows(path: Path) -> List[Dict[str, Any]]:
    episodes: List[Dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            if not line.strip():
                continue
            episodes.append(json.loads(line))
    return episodes


def _resolve_datasets(root: Path) -> List[Path]:
    meta_path = root / META_DIR / EPISODES_FILE
    if meta_path.exists():
        return [root]
    return [p for p in root.iterdir() if p.is_dir()]


def _get_video_features(info: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    features = info.get("features", {}) if isinstance(info, dict) else {}
    video_features: Dict[str, Dict[str, Any]] = {}
    for name, spec in features.items():
        if isinstance(spec, dict) and spec.get("dtype") == "video":
            video_features[name] = spec
    return video_features


def _validate_video(video_path: Optional[Path], expected_length: Optional[int]) -> List[str]:
    if not video_path or not video_path.exists():
        return ["VIDEO_MISSING"]

    try:
        import av
    except Exception as exc:  # pragma: no cover - runtime dependency
        return [f"VIDEO_DECODE_UNAVAILABLE({exc})"]

    try:
        with av.open(str(video_path)) as container:
            stream = container.streams.video[0]
            frame_count = stream.frames

            if frame_count == 0:
                count = 0
                for _ in container.decode(video=0):
                    count += 1
                    if count > 5:
                        break
                if count == 0:
                    return ["VIDEO_EMPTY"]
                frame_count = count

            if expected_length and expected_length > 0 and frame_count != expected_length:
                return [f"FRAME_MISMATCH(video={frame_count}, meta={expected_length})"]

        return []

    except Exception:
        return ["VIDEO_DECODE_FAIL"]


def _annotate_video_errors(errors: List[str], video_key: Optional[str]) -> List[str]:
    if not video_key:
        return errors
    return [f"{err}(video={video_key})" for err in errors]


def check_lerobot_v2_1_quality(
    root: str | Path,
    min_ratio: float = 0.5,
    max_ratio: float = 4.0,
    max_episodes: Optional[int] = None,
) -> Dict[str, Any]:
    """
    Check episode-level quality for LeRobot v2.1 datasets.

    Args:
        root: Dataset root path, or a parent directory containing datasets.
        min_ratio: Minimum length ratio vs task average before flagging as too short.
        max_ratio: Maximum length ratio vs task average before flagging as too long.
        max_episodes: Optional cap on episodes per dataset.

    Returns:
        Dictionary with invalid episodes and summary statistics.
    """
    root_path = Path(root)
    datasets = _resolve_datasets(root_path)

    invalid_episodes: List[Dict[str, Any]] = []
    task_counter: Dict[str, int] = {}
    dataset_errors: List[Dict[str, str]] = []
    total_episodes = 0

    for dataset in datasets:
        meta_path = dataset / META_DIR / EPISODES_FILE
        if not meta_path.exists():
            continue

        info_path = dataset / META_DIR / INFO_FILE
        if not info_path.exists():
            dataset_errors.append({"dataset": str(dataset), "error": "Missing info.json"})
            continue

        try:
            info = _load_json(info_path)
        except Exception as exc:
            dataset_errors.append({"dataset": str(dataset), "error": f"Failed to read info.json: {exc}"})
            continue

        try:
            episodes = _iter_episode_rows(meta_path)
        except Exception as exc:
            dataset_errors.append({"dataset": str(dataset), "error": f"Failed to read episodes.jsonl: {exc}"})
            continue

        video_features = _get_video_features(info)
        file_manager = FileSystemManager(dataset)

        task_groups: Dict[str, List[Dict[str, Any]]] = {}
        for episode in episodes:
            tasks = episode.get("tasks") if isinstance(episode.get("tasks"), list) else None
            task = tasks[0] if tasks else episode.get("task") or "UNKNOWN"
            task_groups.setdefault(task, []).append(episode)

        task_avg: Dict[str, float] = {}
        for task, eps in task_groups.items():
            lengths = [e.get("length", 0) for e in eps if isinstance(e.get("length"), int) and e.get("length", 0) > 0]
            if lengths:
                task_avg[task] = sum(lengths) / len(lengths)

        checked_count = 0
        for episode in episodes:
            episode_index = episode.get("episode_index")
            if episode_index is None:
                continue

            length = episode.get("length", 0)
            tasks = episode.get("tasks") if isinstance(episode.get("tasks"), list) else None
            task = tasks[0] if tasks else episode.get("task") or "UNKNOWN"

            task_counter[task] = task_counter.get(task, 0) + 1
            total_episodes += 1

            ep_name = f"episode_{episode_index:06d}"
            warn: List[str] = []

            video_paths: Dict[str, Path] = {}
            if video_features:
                video_paths = file_manager.get_episode_file_paths(int(episode_index), video_features)["videos"]
            else:
                candidates = list((dataset / "videos").rglob(f"{ep_name}.mp4"))
                if candidates:
                    video_paths = {"video": candidates[0]}

            if not video_paths:
                warn.append("VIDEO_MISSING")
            else:
                for video_key, video_path in video_paths.items():
                    errors = _validate_video(video_path, length)
                    warn.extend(_annotate_video_errors(errors, video_key if len(video_paths) > 1 else None))

            avg_len = task_avg.get(task)
            ratio = None
            if avg_len and avg_len > 0 and isinstance(length, (int, float)):
                ratio = length / avg_len
                if ratio < min_ratio:
                    warn.append("LENGTH_TOO_SHORT")
                if ratio > max_ratio:
                    warn.append("LENGTH_TOO_LONG")

            if warn:
                invalid_episodes.append(
                    {
                        "dataset_path": str(dataset),
                        "task": task,
                        "episode_index": int(episode_index),
                        "episode_name": ep_name,
                        "length": int(length) if isinstance(length, int) else length,
                        "avg_length": avg_len,
                        "ratio": ratio,
                        "errors": warn,
                    }
                )

            checked_count += 1
            if max_episodes and checked_count >= max_episodes:
                break

    total_invalid = len(invalid_episodes)
    invalid_ratio = total_invalid / total_episodes if total_episodes else 0

    return {
        "datasets": [str(p) for p in datasets],
        "invalid_episodes": invalid_episodes,
        "task_counts": task_counter,
        "total_episodes": total_episodes,
        "total_invalid": total_invalid,
        "invalid_ratio": invalid_ratio,
        "dataset_errors": dataset_errors,
        "min_ratio": min_ratio,
        "max_ratio": max_ratio,
    }


def delete_invalid_episodes(
    root: str | Path,
    report: Optional[Dict[str, Any]] = None,
    min_ratio: float = 0.5,
    max_ratio: float = 4.0,
    max_episodes: Optional[int] = None,
    dry_run: bool = False,
) -> Dict[str, Any]:
    """
    Delete invalid episodes detected by check_lerobot_v2_1_quality.

    Args:
        root: Dataset root path, or a parent directory containing datasets.
        report: Optional quality report. If None, a new report is generated.
        min_ratio: Minimum length ratio vs task average before flagging as too short.
        max_ratio: Maximum length ratio vs task average before flagging as too long.
        max_episodes: Optional cap on episodes per dataset (used when report is None).
        dry_run: If True, only show what would be deleted.

    Returns:
        Dictionary with deletion summary and the quality report used.
    """
    if report is None:
        report = check_lerobot_v2_1_quality(
            root,
            min_ratio=min_ratio,
            max_ratio=max_ratio,
            max_episodes=max_episodes,
        )

    invalid_episodes = report.get("invalid_episodes", [])
    invalid_by_dataset: Dict[str, List[int]] = {}
    for entry in invalid_episodes:
        dataset_path = entry.get("dataset_path")
        episode_index = entry.get("episode_index")
        if dataset_path is None or episode_index is None:
            continue
        invalid_by_dataset.setdefault(dataset_path, []).append(int(episode_index))

    deleted: List[Tuple[str, int]] = []
    failed: List[Tuple[str, int]] = []

    for dataset_path, episode_indices in invalid_by_dataset.items():
        editor = LeRobotDatasetEditor(dataset_path)
        for episode_index in sorted(set(episode_indices), reverse=True):
            success = editor.delete_episode(episode_index, dry_run=dry_run)
            if success:
                deleted.append((dataset_path, episode_index))
            else:
                failed.append((dataset_path, episode_index))

    return {
        "report": report,
        "deleted": deleted,
        "failed": failed,
        "dry_run": dry_run,
    }


__all__ = ["check_lerobot_v2_1_quality", "delete_invalid_episodes"]
