"""LeRobot v2.1 dataset merge operations."""

from __future__ import annotations

import contextlib
import json
import shutil
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
import pandas as pd


class DatasetOperationsV21:
    """Dataset operations focused on LeRobot v2.1 merge workflows."""

    def __init__(self, dataset_path: Path):
        self.dataset_path = Path(dataset_path)

    def merge_datasets(
        self,
        source_datasets: List[Path],
        output_path: Path,
        task_mapping: Optional[Dict[str, str]] = None,
        dry_run: bool = False,
    ) -> bool:
        """Merge multiple LeRobot v2.1 datasets into a single dataset."""
        try:
            source_paths = [Path(path) for path in source_datasets]
            output_root = Path(output_path)

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

            if dry_run:
                print(f"\n=== DRY RUN: Merging {len(source_paths)} datasets ===")
                print(f"Output path: {output_root}")
                for i, source_path in enumerate(source_paths, start=1):
                    print(f"Dataset {i}: {source_path}")
                return True

            self._merge_impl(source_paths, output_root, task_mapping=task_mapping)
            return True
        except Exception as e:
            print(f"Error merging datasets with v2.1 operations: {e}")
            return False

    def _load_jsonl(self, file_path: str | Path) -> list[dict[str, Any]]:
        file_path = str(file_path)
        data: list[dict[str, Any]] = []
        if "episodes_stats.jsonl" in file_path:
            try:
                with open(file_path, encoding="utf-8") as f:
                    content = f.read()
                    if content.strip().startswith("[") and content.strip().endswith("]"):
                        return json.loads(content)
                    with contextlib.suppress(json.JSONDecodeError):
                        return json.loads("[" + content + "]")
            except Exception:
                pass

        with open(file_path, encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    with contextlib.suppress(json.JSONDecodeError):
                        data.append(json.loads(line))
        return data

    def _save_jsonl(self, data: list[dict[str, Any]], file_path: str | Path) -> None:
        with open(file_path, "w", encoding="utf-8") as f:
            for item in data:
                f.write(json.dumps(item, ensure_ascii=False) + "\n")

    def _find_episode_parquet(self, dataset_root: str | Path, episode_index: int) -> Path:
        dataset_root = Path(dataset_root)
        expected = dataset_root / "data" / f"chunk-{episode_index // 1000:03d}" / f"episode_{episode_index:06d}.parquet"
        if expected.exists():
            return expected

        matches = sorted((dataset_root / "data").rglob(f"episode_{episode_index:06d}.parquet"))
        if matches:
            return matches[0]

        raise FileNotFoundError(f"Parquet not found for episode {episode_index} in {dataset_root}")

    def _find_video_path(
        self,
        dataset_root: str | Path,
        video_path_template: str,
        chunks_size: int,
        video_key: str,
        episode_index: int,
    ) -> Path:
        dataset_root = Path(dataset_root)
        episode_chunk = episode_index // chunks_size
        candidate = dataset_root / video_path_template.format(
            episode_chunk=episode_chunk,
            video_key=video_key,
            episode_index=episode_index,
        )
        if candidate.exists():
            return candidate

        matches = sorted((dataset_root / "videos").rglob(f"episode_{episode_index:06d}.mp4"))
        matches = [m for m in matches if video_key in str(m.parent)]
        if matches:
            return matches[0]

        raise FileNotFoundError(
            f"Video not found for key={video_key}, episode={episode_index} in {dataset_root}"
        )

    def _merge_stats(self, stats_list: list[dict[str, Any]]) -> dict[str, Any]:
        if not stats_list:
            return {}

        merged = {}
        common_features = set(stats_list[0].keys())
        for stats in stats_list[1:]:
            common_features &= set(stats.keys())

        for feature in stats_list[0]:
            if feature not in common_features:
                continue
            merged[feature] = {}
            for stat_type in ["mean", "std", "min", "max", "count"]:
                if not all(stat_type in stats[feature] for stats in stats_list):
                    continue

                values = [stats[feature][stat_type] for stats in stats_list]
                try:
                    if stat_type == "count":
                        merged[feature][stat_type] = [sum(v[0] for v in values)]
                    elif stat_type == "mean":
                        if all("count" in stats[feature] for stats in stats_list):
                            counts = [stats[feature]["count"][0] for stats in stats_list]
                            total = sum(counts)
                            weighted = [np.asarray(v) * c / total for v, c in zip(values, counts)]
                            merged[feature][stat_type] = np.sum(weighted, axis=0).tolist()
                        else:
                            merged[feature][stat_type] = np.mean(np.asarray(values), axis=0).tolist()
                    elif stat_type == "std":
                        if all("count" in stats[feature] for stats in stats_list):
                            counts = [stats[feature]["count"][0] for stats in stats_list]
                            total = sum(counts)
                            variances = [np.asarray(v) ** 2 for v in values]
                            weighted = [v * c / total for v, c in zip(variances, counts)]
                            merged[feature][stat_type] = np.sqrt(np.sum(weighted, axis=0)).tolist()
                        else:
                            merged[feature][stat_type] = np.mean(np.asarray(values), axis=0).tolist()
                    elif stat_type == "min":
                        merged[feature][stat_type] = np.minimum.reduce(np.asarray(values)).tolist()
                    elif stat_type == "max":
                        merged[feature][stat_type] = np.maximum.reduce(np.asarray(values)).tolist()
                except Exception:
                    merged[feature][stat_type] = values[0]
        return merged

    def _copy_data_files(
        self,
        output_folder: Path,
        episode_mapping: list[tuple[str, int, int]],
        episode_to_frame_index: dict[int, int],
        folder_task_mapping: dict[str, dict[int, int]],
        chunks_size: int,
    ) -> None:
        for old_folder, old_index, new_index in episode_mapping:
            source_path = self._find_episode_parquet(old_folder, old_index)
            df = pd.read_parquet(source_path)

            if "episode_index" in df.columns:
                df["episode_index"] = new_index
            if "index" in df.columns:
                first_index = episode_to_frame_index[new_index]
                df["index"] = [first_index + i for i in range(len(df))]
            if "task_index" in df.columns and old_folder in folder_task_mapping:
                current_task_index = int(df["task_index"].iloc[0])
                if current_task_index in folder_task_mapping[old_folder]:
                    df["task_index"] = folder_task_mapping[old_folder][current_task_index]

            chunk_index = new_index // chunks_size
            chunk_dir = output_folder / "data" / f"chunk-{chunk_index:03d}"
            chunk_dir.mkdir(parents=True, exist_ok=True)
            dest_path = chunk_dir / f"episode_{new_index:06d}.parquet"
            df.to_parquet(dest_path, index=False)

    def _copy_videos(
        self,
        output_folder: Path,
        episode_mapping: list[tuple[str, int, int]],
        info: dict[str, Any],
    ) -> None:
        chunks_size = info["chunks_size"]
        video_path_template = info["video_path"]
        video_keys = [
            feature_name
            for feature_name, feature_info in info["features"].items()
            if feature_info.get("dtype") == "video"
        ]

        for old_folder, old_index, new_index in episode_mapping:
            new_episode_chunk = new_index // chunks_size
            for video_key in video_keys:
                source_video_path = self._find_video_path(
                    old_folder,
                    video_path_template,
                    chunks_size,
                    video_key,
                    old_index,
                )
                dest_video_path = output_folder / video_path_template.format(
                    episode_chunk=new_episode_chunk,
                    video_key=video_key,
                    episode_index=new_index,
                )
                dest_video_path.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(source_video_path, dest_video_path)

    def _merge_impl(
        self,
        source_folders: list[Path],
        output_root: Path,
        task_mapping: Optional[Dict[str, str]] = None,
    ) -> None:
        if output_root.exists():
            raise FileExistsError(f"Output already exists: {output_root}")

        (output_root / "meta").mkdir(parents=True, exist_ok=True)

        all_episodes = []
        all_episodes_stats = []
        all_tasks = []
        episode_mapping: list[tuple[str, int, int]] = []
        episode_to_frame_index: dict[int, int] = {}
        folder_task_mapping: dict[str, dict[int, int]] = {}
        task_desc_to_new_index: dict[str, int] = {}

        total_frames = 0
        total_episodes = 0
        cumulative_frame_count = 0
        total_videos = 0

        with open(source_folders[0] / "meta" / "info.json", encoding="utf-8") as f:
            info = json.load(f)
        chunks_size = info.get("chunks_size", 1000)

        for folder_path in source_folders:
            folder = str(folder_path)
            with open(folder_path / "meta" / "info.json", encoding="utf-8") as f:
                folder_info = json.load(f)
            total_videos += folder_info.get("total_videos", 0)

            episodes = self._load_jsonl(folder_path / "meta" / "episodes.jsonl")
            episodes_stats = []
            episodes_stats_path = folder_path / "meta" / "episodes_stats.jsonl"
            if episodes_stats_path.exists():
                episodes_stats = self._load_jsonl(episodes_stats_path)
            stats_map = {item["episode_index"]: item for item in episodes_stats if "episode_index" in item}

            folder_tasks = self._load_jsonl(folder_path / "meta" / "tasks.jsonl")
            folder_task_mapping[folder] = {}
            for task in folder_tasks:
                task_desc = task.get("task", "")
                if task_mapping:
                    task_desc = task_mapping.get(task_desc, task_desc)
                old_task_index = task["task_index"]
                if task_desc not in task_desc_to_new_index:
                    new_task_index = len(all_tasks)
                    task_desc_to_new_index[task_desc] = new_task_index
                    all_tasks.append({"task_index": new_task_index, "task": task_desc})
                folder_task_mapping[folder][old_task_index] = task_desc_to_new_index[task_desc]

            for episode in episodes:
                old_index = int(episode["episode_index"])
                new_index = total_episodes

                new_episode = dict(episode)
                new_episode["episode_index"] = new_index
                all_episodes.append(new_episode)

                if old_index in stats_map:
                    stat = dict(stats_map[old_index])
                    stat["episode_index"] = new_index
                    all_episodes_stats.append(stat)

                episode_mapping.append((folder, old_index, new_index))
                episode_to_frame_index[new_index] = cumulative_frame_count
                cumulative_frame_count += int(episode["length"])
                total_frames += int(episode["length"])
                total_episodes += 1

        self._save_jsonl(all_episodes, output_root / "meta" / "episodes.jsonl")
        self._save_jsonl(all_episodes_stats, output_root / "meta" / "episodes_stats.jsonl")
        self._save_jsonl(all_tasks, output_root / "meta" / "tasks.jsonl")

        stats_list = []
        for folder_path in source_folders:
            stats_path = folder_path / "meta" / "stats.json"
            if stats_path.exists():
                with open(stats_path, encoding="utf-8") as f:
                    stats_list.append(json.load(f))
        if stats_list:
            merged_stats = self._merge_stats(stats_list)
            with open(output_root / "meta" / "stats.json", "w", encoding="utf-8") as f:
                json.dump(merged_stats, f, indent=2, ensure_ascii=False)

        merged_info = dict(info)
        merged_info["total_episodes"] = total_episodes
        merged_info["total_frames"] = total_frames
        merged_info["total_tasks"] = len(all_tasks)
        merged_info["total_chunks"] = (total_episodes + chunks_size - 1) // chunks_size
        merged_info["total_videos"] = total_videos
        merged_info["splits"] = {"train": f"0:{total_episodes}"}
        with open(output_root / "meta" / "info.json", "w", encoding="utf-8") as f:
            json.dump(merged_info, f, indent=2, ensure_ascii=False)

        self._copy_videos(output_root, episode_mapping, merged_info)
        self._copy_data_files(
            output_root,
            episode_mapping,
            episode_to_frame_index,
            folder_task_mapping,
            chunks_size,
        )

        image_keys = [k for k in merged_info.get("features", {}) if k.startswith("observation.images.")]
        if image_keys:
            images_dir = output_root / "images"
            images_dir.mkdir(parents=True, exist_ok=True)
            for image_key in image_keys:
                (images_dir / image_key).mkdir(parents=True, exist_ok=True)

        print(f"Merged {total_episodes} episodes with {total_frames} frames into {output_root}")
