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

Single-task mode (--dataset-root points at a data/ directory):

    python make_robotwin_meta.py \\
        --dataset-root .../place_can_basket/aloha-agilex_clean_50/data \\
        --output meta/place_can_basket.json

Multi-task mode (--dataset-root points at the top-level directory that
contains one subdirectory per task):

    python make_robotwin_meta.py \\
        --dataset-root /work/xuan/dataset/robotwin_dataset \\
        --output meta/robotwin_all.json \\
        --data-subdir aloha-agilex_clean_50

Instructions are auto-detected from sibling 'instructions/' directories.
"""

import argparse
import json
from pathlib import Path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate X-VLA metadata for RoboTwin HDF5 episodes.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "--dataset-root", required=True,
        help="Top-level directory containing task folders (multi-task), "
             "or a single data/ directory (single-task).",
    )
    parser.add_argument("--output", required=True,
                        help="Where to write the X-VLA metadata JSON.")
    parser.add_argument(
        "--dataset-name",
        default="robotwin2_clean",
        choices=["robotwin2_clean", "robotwin2_abs_ee"],
        help="Dataset key registered in X-VLA (default: robotwin2_clean).",
    )
    parser.add_argument(
        "--data-subdir", default=None,
        help="Subdirectory name that identifies each task's dataset "
             "(e.g. 'aloha-agilex_clean_50'). When set, enables multi-task "
             "discovery: the script looks for <task>/<data-subdir>/data/ "
             "under --dataset-root.",
    )
    parser.add_argument(
        "--glob", default="*.hdf5",
        help="Glob pattern for episode files (default: *.hdf5).",
    )
    parser.add_argument(
        "--observation-key", nargs="+",
        default=[
            "observation/countertop_camera/rgb",
            "observation/left_camera/rgb",
            "observation/right_camera/rgb",
        ],
        help="Image dataset keys inside each HDF5 file.",
    )

    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 discover_tasks(root: Path, data_subdir: str) -> list[tuple[str, Path, Path]]:
    """Find all <task>/<data_subdir>/{data/,instructions/} under root.

    Returns list of (task_name, data_dir, instruction_dir).
    """
    tasks = []
    for task_dir in sorted(root.iterdir()):
        if not task_dir.is_dir():
            continue
        base = task_dir / data_subdir
        data_dir = base / "data"
        instr_dir = base / "instructions"
        if data_dir.is_dir():
            tasks.append((task_dir.name, data_dir, instr_dir if instr_dir.is_dir() else None))
    return tasks


def build_instruction_maps(
    episodes: list[tuple[str, Path | None]],
    split: str = "seen",
) -> tuple[dict[str, str], dict[str, list[str]]]:
    """Build instruction_map and lang_aug_map from per-episode JSON files.

    Args:
        episodes: list of (hdf5_absolute_path, instruction_json_path_or_None)
        split: "seen" or "unseen"

    Returns:
        instruction_map: {hdf5_absolute_path: canonical_instruction}
        lang_aug_map:    {canonical_instruction: [all paraphrases]}
    """
    instruction_map: dict[str, str] = {}
    lang_aug_map: dict[str, list[str]] = {}
    missing = []

    for hdf5_path, json_path in episodes:
        if json_path is None or not json_path.exists():
            missing.append(hdf5_path)
            continue

        with open(json_path, "r", encoding="utf-8") as f:
            instr_data = json.load(f)

        if split not in instr_data:
            available = list(instr_data.keys())
            raise KeyError(
                f"Split '{split}' not found in {json_path}. Available: {available}"
            )

        instructions = instr_data[split]
        if not instructions:
            raise ValueError(
                f"Empty instruction list for split '{split}' in {json_path}"
            )

        canonical = instructions[0]
        instruction_map[hdf5_path] = canonical

        if len(instructions) > 1:
            lang_aug_map[canonical] = instructions

    if missing:
        raise FileNotFoundError(
            f"Missing instruction JSONs for {len(missing)} episodes, "
            f"e.g.: {missing[:3]}"
        )

    return instruction_map, lang_aug_map


def collect_single_task(
    data_dir: Path, glob_pattern: str,
) -> tuple[list[str], Path | None]:
    """Collect episodes from a single data directory.

    Returns (datalist, instruction_dir_or_None).
    """
    datalist = sorted(str(p.resolve()) for p in data_dir.rglob(glob_pattern))
    instr_dir = None
    for candidate in [data_dir / "instructions", data_dir.parent / "instructions"]:
        if candidate.is_dir():
            instr_dir = candidate
            break
    return datalist, instr_dir


def pair_episodes_with_instructions(
    datalist: list[str], instr_dir: Path | None,
) -> list[tuple[str, Path | None]]:
    """Pair each HDF5 path with its corresponding instruction JSON."""
    pairs = []
    for hdf5_path in datalist:
        json_path = None
        if instr_dir is not None:
            json_path = instr_dir / f"{Path(hdf5_path).stem}.json"
        pairs.append((hdf5_path, json_path))
    return pairs


def main() -> None:
    args = parse_args()
    dataset_root = Path(args.dataset_root).expanduser().resolve()
    output = Path(args.output).expanduser().resolve()

    all_datalist: list[str] = []
    all_episode_pairs: list[tuple[str, Path | None]] = []
    has_instructions = False

    if args.data_subdir:
        # --- Multi-task discovery ---
        tasks = discover_tasks(dataset_root, args.data_subdir)
        if not tasks:
            raise FileNotFoundError(
                f"No task directories with '{args.data_subdir}/data/' "
                f"found under {dataset_root}"
            )

        for task_name, data_dir, instr_dir in tasks:
            datalist = sorted(
                str(p.resolve()) for p in data_dir.rglob(args.glob)
            )
            if not datalist:
                print(f"  WARNING: no HDF5 files in {data_dir}, skipping {task_name}")
                continue

            pairs = pair_episodes_with_instructions(datalist, instr_dir)
            all_datalist.extend(datalist)
            all_episode_pairs.extend(pairs)

            if instr_dir is not None:
                has_instructions = True

            print(f"  [{task_name}] {len(datalist)} episodes"
                  f"{' + instructions' if instr_dir else ''}")

        print(f"Discovered {len(tasks)} tasks, "
              f"{len(all_datalist)} total episodes")
    else:
        # --- Single-task mode ---
        datalist, instr_dir = collect_single_task(dataset_root, args.glob)
        if not datalist:
            raise FileNotFoundError(
                f"No files matched {args.glob!r} under {dataset_root}"
            )
        all_datalist = datalist
        all_episode_pairs = pair_episodes_with_instructions(datalist, instr_dir)
        has_instructions = instr_dir is not None

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

    # --- Resolve instructions ---
    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 has_instructions:
        print(f"Building instruction maps (split={args.instruction_split!r})...")
        instruction_map, lang_aug_map = build_instruction_maps(
            all_episode_pairs, split=args.instruction_split,
        )
        meta["instruction_map"] = instruction_map
        if lang_aug_map:
            meta["lang_aug_map"] = lang_aug_map
    else:
        raise ValueError(
            "No instruction source found. Use one of:\n"
            "  --language-instruction-key  (HDF5 key)\n"
            "  --default-instruction       (single string)\n"
            "  --instruction-map-json      (pre-built JSON)\n"
            "Or ensure 'instructions/' directories exist alongside data/."
        )

    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")

    n_episodes = len(all_datalist)
    n_instr = len(meta.get("instruction_map", {}))
    n_aug = len(meta.get("lang_aug_map", {}))
    print(f"\nWrote {n_episodes} episodes to {output}")
    if n_instr:
        print(f"  instruction_map: {n_instr} entries (keyed by absolute path)")
    if n_aug:
        print(f"  lang_aug_map:    {n_aug} entries (for training augmentation)")


if __name__ == "__main__":
    main()