#!/usr/bin/env python3
"""Generate X-VLA metadata for RoboReal HDF5 episodes.

Expected dataset layout:

    <dataset-root>/
      <task-name>/
        data/
          episode0.hdf5
          ...

Or:

    <dataset-root>/
      <task-name>/
        <variant>/
          data/
            episode0.hdf5
            ...

Or:

    <dataset-root>/
      <domain>/
        <task-name>/
          <variant>/
            data/
              episode0.hdf5
              ...

Examples:

    /shared_work/xuan/dataset/roboreal_data/study/move_cup/clean/data/episode0.hdf5
    /shared_work/xuan/dataset/roboreal_data/office/move_items_around/office_d10/data/episode0.hdf5

The generated metadata aggregates all matching episode files into a single
X-VLA JSON so training can consume either a study subtree or the whole
RoboReal tree at once.
"""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path

from make_robotwin_meta import build_instruction_maps, pair_episodes_with_instructions


DEFAULT_OBSERVATION_KEYS = [
    "observation/countertop_camera/rgb",
    "observation/left_camera/rgb",
    "observation/right_camera/rgb",
    "observation/head_camera/rgb",
]
CAMERA_PRIORITY = ["countertop_camera", "left_camera", "right_camera", "head_camera"]

DEFAULT_IGNORED_DIRS = [".cache", ".git", "__pycache__"]
HDF5_SIGNATURE = b"\x89HDF\r\n\x1a\n"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate X-VLA metadata for RoboReal datasets.",
    )
    parser.add_argument(
        "--dataset-root",
        required=True,
        help="Root RoboReal directory containing task/data, task/variant/data, or domain/task/variant/data leaves.",
    )
    parser.add_argument(
        "--output",
        required=True,
        help="Path to the output metadata JSON.",
    )
    parser.add_argument(
        "--dataset-name",
        default="robotwin2_clean",
        choices=["robotwin2_clean", "robotwin2_abs_ee"],
        help="Dataset key registered in X-VLA.",
    )
    parser.add_argument(
        "--glob",
        default="*.hdf5",
        help="Glob pattern for episode files inside each data directory.",
    )
    parser.add_argument(
        "--observation-key",
        nargs="+",
        default=None,
        help="Image dataset keys inside each HDF5 file. If omitted, infer from the first episode and prefer head/left/right cameras when present.",
    )
    parser.add_argument(
        "--ignore-dir",
        nargs="+",
        default=DEFAULT_IGNORED_DIRS,
        help="Directory names to skip anywhere under dataset root.",
    )
    parser.add_argument(
        "--variant",
        nargs="*",
        default=None,
        help="Optional allowlist of dataset variants to include, e.g. clean study_d6.",
    )
    instr = parser.add_argument_group("instruction sources")
    instr.add_argument(
        "--language-instruction-key",
        help="HDF5 dataset key for the instruction.",
    )
    instr.add_argument(
        "--default-instruction",
        help="Single instruction shared by every episode.",
    )
    instr.add_argument(
        "--instruction-map-json",
        help="Pre-built JSON file mapping episode path/name/stem to instruction text.",
    )
    instr.add_argument(
        "--instruction-split",
        default="seen",
        choices=["seen", "unseen"],
        help="Which split to use from per-episode instruction JSONs (default: seen).",
    )
    return parser.parse_args()


def task_to_instruction(task_name: str) -> str:
    base = task_name
    for suffix in ("_clean", "_demo_clean", "_demo_randomized", "_randomized"):
        if base.endswith(suffix):
            base = base[: -len(suffix)]
            break

    text = base.replace("_", " ").strip()
    text = re.sub(r"\s+", " ", text)
    if not text:
        raise ValueError(f"Could not derive instruction from task name: {task_name}")
    return text[0].upper() + text[1:] + "."


def should_skip(path: Path, ignored_dirs: set[str]) -> bool:
    return any(part in ignored_dirs for part in path.parts)


def is_valid_hdf5(path: Path) -> bool:
    try:
        with open(path, "rb") as f:
            return f.read(len(HDF5_SIGNATURE)) == HDF5_SIGNATURE
    except OSError:
        return False


def infer_observation_keys(representative_hdf5s: list[str]) -> list[str]:
    """Return camera keys present in ALL given HDF5 files (intersection)."""
    import h5py

    common: set[str] | None = None
    for path in representative_hdf5s:
        with h5py.File(path, "r") as f:
            if "observation" not in f:
                raise KeyError(f"'observation' group not found in {path}")
            available = {
                camera_name
                for camera_name, camera_group in f["observation"].items()
                if isinstance(camera_group, h5py.Group) and "rgb" in camera_group
            }
        common = available if common is None else common & available

    if not common:
        raise ValueError(
            f"No observation/*/rgb datasets common across all representative episodes"
        )

    ordered = [name for name in CAMERA_PRIORITY if name in common]
    ordered.extend(sorted(name for name in common if name not in ordered))
    return [f"observation/{camera_name}/rgb" for camera_name in ordered]


def find_instruction_dir(data_dir: Path) -> Path | None:
    for candidate in (data_dir / "instructions", data_dir.parent / "instructions"):
        if candidate.is_dir():
            return candidate
    return None


def iter_data_dirs(dataset_root: Path, ignored_dirs: set[str]) -> list[tuple[Path, str]]:
    data_dirs: list[tuple[Path, str]] = []
    for data_dir in sorted(dataset_root.rglob("data")):
        if not data_dir.is_dir():
            continue
        if should_skip(data_dir, ignored_dirs):
            continue
        relative = data_dir.relative_to(dataset_root)
        layout = parse_layout(relative)
        if layout is None:
            continue
        data_dirs.append((data_dir, layout))
    return data_dirs


def parse_layout(relative: Path) -> str | None:
    parts = relative.parts
    if not parts or parts[-1] != "data":
        return None
    if len(parts) == 2:
        return "task"
    if len(parts) == 3:
        return "task_variant"
    if len(parts) == 4:
        return "domain_task_variant"
    return None


def collect_datasets(
    dataset_root: Path,
    glob_pattern: str,
    ignored_dirs: set[str],
    variants: set[str] | None,
) -> tuple[
    list[str],
    list[str],
    list[tuple[str, Path | None]],
    dict[str, str],
    dict[str, int],
    dict[str, int],
    int,
]:
    datalist: list[str] = []
    representatives: list[str] = []
    episode_pairs: list[tuple[str, Path | None]] = []
    instruction_map: dict[str, str] = {}
    dataset_counts: dict[str, int] = {}
    skipped_counts: dict[str, int] = {}
    datasets_with_instruction_dir = 0

    for data_dir, layout in iter_data_dirs(dataset_root, ignored_dirs):
        relative = data_dir.relative_to(dataset_root)
        if layout == "task":
            domain_name = dataset_root.name
            task_name, leaf_name = relative.parts
            variant_name = "default"
        elif layout == "task_variant":
            domain_name = dataset_root.name
            task_name, variant_name, leaf_name = relative.parts
        else:
            domain_name, task_name, variant_name, leaf_name = relative.parts

        if variants is not None and variant_name not in variants:
            continue

        raw_episodes = sorted(p.resolve() for p in data_dir.rglob(glob_pattern))
        episodes = [str(p) for p in raw_episodes if is_valid_hdf5(p)]
        skipped = len(raw_episodes) - len(episodes)
        if skipped:
            dataset_key = f"{domain_name}/{task_name}/{variant_name}"
            skipped_counts[dataset_key] = skipped
        if not episodes:
            continue

        instruction = task_to_instruction(task_name)
        dataset_key = f"{domain_name}/{task_name}/{variant_name}"
        dataset_counts[dataset_key] = len(episodes)
        datalist.extend(episodes)
        representatives.append(episodes[0])
        instr_dir = find_instruction_dir(data_dir)
        if instr_dir is not None:
            datasets_with_instruction_dir += 1
        episode_pairs.extend(pair_episodes_with_instructions(episodes, instr_dir))
        for episode in episodes:
            instruction_map[episode] = instruction

    return (
        datalist,
        representatives,
        episode_pairs,
        instruction_map,
        dataset_counts,
        skipped_counts,
        datasets_with_instruction_dir,
    )


def main() -> None:
    args = parse_args()
    dataset_root = Path(args.dataset_root).expanduser().resolve()
    output = Path(args.output).expanduser().resolve()
    variants = set(args.variant) if args.variant else None
    ignored_dirs = set(args.ignore_dir)

    (
        datalist,
        representatives,
        episode_pairs,
        derived_instruction_map,
        dataset_counts,
        skipped_counts,
        datasets_with_instruction_dir,
    ) = collect_datasets(
        dataset_root=dataset_root,
        glob_pattern=args.glob,
        ignored_dirs=ignored_dirs,
        variants=variants,
    )
    if not datalist:
        raise FileNotFoundError(
            f"No files matched {args.glob!r} under {dataset_root}"
        )

    observation_keys = args.observation_key or infer_observation_keys(representatives)

    meta = {
        "dataset_name": args.dataset_name,
        "observation_key": observation_keys,
        "datalist": datalist,
    }

    if args.language_instruction_key:
        meta["language_instruction_key"] = args.language_instruction_key
    elif args.default_instruction:
        meta["default_instruction"] = args.default_instruction
    elif args.instruction_map_json:
        with open(args.instruction_map_json, "r", encoding="utf-8") as f:
            meta["instruction_map"] = json.load(f)
    elif datasets_with_instruction_dir == len(dataset_counts):
        print(f"Building instruction maps (split={args.instruction_split!r})...")
        instruction_map, lang_aug_map = build_instruction_maps(
            episode_pairs,
            split=args.instruction_split,
        )
        meta["instruction_map"] = instruction_map
        if lang_aug_map:
            meta["lang_aug_map"] = lang_aug_map
    elif datasets_with_instruction_dir:
        raise ValueError(
            "Found instructions/ directories for only part of the dataset tree. "
            "Use an explicit instruction source, or make instruction files "
            "available for every included dataset."
        )
    else:
        meta["instruction_map"] = derived_instruction_map

    output.parent.mkdir(parents=True, exist_ok=True)
    with open(output, "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)
        f.write("\n")

    print(f"Wrote {len(datalist)} episodes to {output}")
    print(f"Datasets: {len(dataset_counts)}")
    print(f"Observation keys: {observation_keys}")
    if "language_instruction_key" in meta:
        print(f"Instruction source: HDF5 key {meta['language_instruction_key']}")
    elif "default_instruction" in meta:
        print("Instruction source: shared default instruction")
    elif args.instruction_map_json:
        print(f"Instruction source: JSON map {args.instruction_map_json}")
    elif "lang_aug_map" in meta:
        print(f"Instruction source: per-episode instruction JSONs ({len(meta['lang_aug_map'])} augmented entries)")
    else:
        print("Instruction source: task-name-derived fallback")
    for dataset_key, count in dataset_counts.items():
        print(f"  {dataset_key}: {count}")
    if skipped_counts:
        print("Skipped non-HDF5 files:")
        for dataset_key, count in skipped_counts.items():
            print(f"  {dataset_key}: {count}")


if __name__ == "__main__":
    main()