#!/usr/bin/env python3
"""RoboTwin 数据准备脚本：将原始 ZIP 格式转换为 RobotWinDataset 期望的 qpos+videos+metas 格式。

用法:
    python scripts/data/prepare_robotwin.py \\
        --input /root/autol-tmp/data/robotwin3/dataset \\
        --output /root/autol-tmp/data/robotwin_gear

RoboTwin 原始结构 (ZIP):
    <task_dir>/
    ├── franka_clean_50.zip
    │   └── franka_clean_50/
    │       ├── scene_info.json
    │       ├── instructions/episodeN.json   # {"seen": [...], "unseen": [...]}
    │       ├── video/episodeN.mp4
    │       └── _traj_data/episodeN.pkl      # {arm_key: [seg1, seg2, ...]}
    ├── aloha-agilex_clean_50.zip
    └── ...

输出结构:
    <output_dir>/
    ├── <task>_<robot>/
    │   ├── qpos/
    │   │   ├── episode0.pt    # [T_qpos, 14] float32
    │   │   └── episode1.pt
    │   ├── videos/
    │   │   ├── episode0.mp4
    │   │   └── episode1.mp4
    │   └── metas/
    │       ├── task_0.txt
    │       └── task_1.txt
    └── ...
"""

import os, sys, json, argparse, logging, pickle, shutil
from pathlib import Path

import numpy as np
import torch
import zipfile

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

# 14-dim layout: 前 7 = right arm, 后 7 = left arm
# 单臂机器人：使用的手臂放前 N 位，其余补 0
ROBOT_CONFIG = {
    "franka":        {"arm": "right", "dims": 7},
    "aloha-agilex":  {"arm": "left",  "dims": 6},
    "arx-x5":        {"arm": "left",  "dims": 6},
    "ur5":           {"arm": "left",  "dims": 6},
    "piper":         {"arm": "left",  "dims": 6},
}

OUTPUT_STATE_DIM = 14


def get_robot_name(zip_path: Path) -> str | None:
    """从 ZIP 文件名提取机器人名称（去掉 _clean_50.zip 后缀）。"""
    name = zip_path.stem  # e.g. franka_clean_50
    # Remove _clean_50 or _50 suffix
    for suffix in ["_clean_50", "_50"]:
        if name.endswith(suffix):
            name = name[:-len(suffix)]
            break
    return name


def get_arm_key(robot: str) -> str:
    """根据机器人类型返回 pkl 中的 arm key。"""
    cfg = ROBOT_CONFIG.get(robot)
    if cfg is None:
        logger.warning(f"Unknown robot {robot}, trying right_joint_path")
        return "right_joint_path"
    return f"{cfg['arm']}_joint_path"


def merge_trajectory_segments(segments: list) -> np.ndarray:
    """合并多段轨迹为一个连续数组。

    Args:
        segments: list of dict, each with "position" key [T_i, D]

    Returns:
        concatenated position array [sum(T_i), D]
    """
    arrays = [seg["position"] for seg in segments if seg.get("position") is not None]
    if not arrays:
        return np.zeros((0, arrays[0].shape[1])) if arrays else np.zeros((0, 1))
    return np.concatenate(arrays, axis=0)


def pad_to_14dim(arr: np.ndarray, joint_dim: int) -> np.ndarray:
    """Pad joint positions to 14-dim。

    布局：前 7 = right arm, 后 7 = left arm。
    - franka (7, right): 放在前 7 维
    - aloha (6, left): 放在后 6 维（前补 0）
    """
    if arr.ndim == 1:
        padded = np.zeros(OUTPUT_STATE_DIM, dtype=np.float32)
        if joint_dim == 7:
            padded[:joint_dim] = arr.astype(np.float32)
        else:  # 6-dim → 后 6 维
            padded[OUTPUT_STATE_DIM - joint_dim:] = arr.astype(np.float32)
    else:
        padded = np.zeros((arr.shape[0], OUTPUT_STATE_DIM), dtype=np.float32)
        if joint_dim == 7:
            padded[:, :joint_dim] = arr.astype(np.float32)
        else:  # 6-dim → 后 6 维
            padded[:, OUTPUT_STATE_DIM - joint_dim:] = arr.astype(np.float32)
    return padded


def get_task_description(zf: zipfile.ZipFile, robot_prefix: str, ep_idx: int) -> str:
    """从 instruction JSON 读取任务描述。

    Returns:
        第一条 "seen" 指令，或 fallback 文本
    """
    try:
        inst_path = f"{robot_prefix}/instructions/episode{ep_idx}.json"
        with zf.open(inst_path) as f:
            inst = json.load(f)
        seen = inst.get("seen", [])
        if seen:
            return seen[0]
        unseen = inst.get("unseen", [])
        if unseen:
            return unseen[0]
    except Exception as e:
        logger.debug(f"  Cannot read instruction: {e}")
    return "Manipulate the object on the table"


def process_robot_zip(
    zip_path: Path,
    output_dir: Path,
    delete_after: bool = False,
) -> int:
    """处理一个 RoboTwin ZIP 文件。

    Returns:
        成功处理的 episode 数
    """
    robot = get_robot_name(zip_path)
    if robot is None:
        logger.warning(f"Cannot determine robot from {zip_path.name}, skipping")
        return 0

    # 输出目录：{task}_{robot}
    task_name = zip_path.parent.name
    out_subdir = output_dir / f"{task_name}_{robot}"
    qpos_dir = out_subdir / "qpos"
    video_dir = out_subdir / "videos"
    meta_dir = out_subdir / "metas"
    qpos_dir.mkdir(parents=True, exist_ok=True)
    video_dir.mkdir(parents=True, exist_ok=True)
    meta_dir.mkdir(parents=True, exist_ok=True)

    cfg = ROBOT_CONFIG.get(robot)
    if cfg is None:
        logger.warning(f"Unknown robot {robot}, skipping {zip_path}")
        return 0

    joint_dim = cfg["dims"]
    arm_key = get_arm_key(robot)
    robot_prefix = f"{robot}_clean_50"  # ZIP 内部目录名

    try:
        zf = zipfile.ZipFile(str(zip_path))
    except Exception as e:
        logger.error(f"Cannot open {zip_path}: {e}")
        return 0

    # 从 scene_info.json 获取 episode 列表
    try:
        with zf.open(f"{robot_prefix}/scene_info.json") as f:
            scene_info = json.load(f)
    except Exception as e:
        logger.warning(f"No scene_info in {zip_path}: {e}")
        zf.close()
        return 0

    # 列出所有 episode
    episode_keys = sorted([k for k in scene_info if k.startswith("episode_")],
                          key=lambda x: int(x.split("_")[1]))
    if not episode_keys:
        logger.warning(f"No episodes in scene_info of {zip_path}")
        zf.close()
        return 0

    success_count = 0
    for ep_key in episode_keys:
        ep_idx = int(ep_key.split("_")[1])
        ep_name = f"episode{ep_idx}"

        # 检查是否有轨迹数据
        traj_path = f"{robot_prefix}/_traj_data/{ep_name}.pkl"
        video_path_in = f"{robot_prefix}/video/{ep_name}.mp4"

        if traj_path not in zf.namelist():
            logger.debug(f"  No traj data for {ep_name} in {zip_path.name}")
            continue
        if video_path_in not in zf.namelist():
            logger.debug(f"  No video for {ep_name} in {zip_path.name}")
            continue

        try:
            # --- 轨迹处理 ---
            with zf.open(traj_path) as f:
                traj_data = pickle.load(f)

            segments = traj_data.get(arm_key, [])
            if not segments:
                logger.debug(f"  No {arm_key} segments for {ep_name}")
                continue

            # 合并多段轨迹
            pos = merge_trajectory_segments(segments)  # [T, joint_dim]
            if pos.shape[0] < 24:  # 至少需要 num_frames + action_horizon = 24
                logger.debug(f"  Too few frames ({pos.shape[0]}) for {ep_name}")
                continue

            # Pad 到 14-dim
            pos_14 = pad_to_14dim(pos, joint_dim)  # [T, 14]

            # 保存为 .pt
            qpos_path = qpos_dir / f"{ep_name}.pt"
            torch.save(torch.from_numpy(pos_14), qpos_path)

            # --- 视频提取 ---
            video_out_path = video_dir / f"{ep_name}.mp4"
            with zf.open(video_path_in) as src, open(video_out_path, "wb") as dst:
                shutil.copyfileobj(src, dst)

            # --- 任务描述 ---
            task_desc = get_task_description(zf, robot_prefix, ep_idx)
            meta_path = meta_dir / f"task_{ep_idx}.txt"
            with open(meta_path, "w") as f:
                f.write(task_desc)

            success_count += 1

        except Exception as e:
            logger.warning(f"  Error processing {ep_name} in {zip_path.name}: {e}")
            continue

    zf.close()

    # 删除 ZIP 释放空间
    if delete_after and success_count > 0:
        try:
            zip_path.unlink()
            logger.info(f"  Deleted {zip_path.name}")
        except Exception as e:
            logger.warning(f"  Cannot delete {zip_path.name}: {e}")

    if success_count > 0:
        logger.info(f"{zip_path.name}: {success_count}/{len(episode_keys)} episodes extracted -> {out_subdir}")

    return success_count


def main():
    parser = argparse.ArgumentParser(description="Prepare RoboTwin data for DreamZero training")
    parser.add_argument("--input", "-i", required=True,
                        help="RoboTwin 数据目录 (含任务子目录)")
    parser.add_argument("--output", "-o", required=True,
                        help="输出目录")
    parser.add_argument("--delete-zip", action="store_true",
                        help="处理完成后删除 ZIP 文件（节省空间）")
    parser.add_argument("--max-tasks", type=int, default=None,
                        help="最多处理前 N 个任务（用于测试）")
    parser.add_argument("--num-workers", type=int, default=4,
                        help="并行处理数（暂未实现）")
    args = parser.parse_args()

    input_dir = Path(args.input)
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 扫描所有任务目录
    task_dirs = sorted([d for d in input_dir.iterdir() if d.is_dir()])
    logger.info(f"Found {len(task_dirs)} task directories")

    if args.max_tasks:
        task_dirs = task_dirs[:args.max_tasks]
        logger.info(f"Limited to {args.max_tasks} tasks")

    total_episodes = 0
    total_zips = 0

    for task_dir in task_dirs:
        # 找到所有 ZIP 文件
        zip_files = sorted(task_dir.glob("*_clean_50.zip"))
        if not zip_files:
            logger.warning(f"No ZIP files in {task_dir}")
            continue

        for zip_path in zip_files:
            try:
                ep_count = process_robot_zip(
                    zip_path, output_dir, delete_after=args.delete_zip
                )
                total_episodes += ep_count
                total_zips += 1
            except Exception as e:
                logger.error(f"Fatal error processing {zip_path}: {e}")
                continue

    logger.info(f"Done! Processed {total_zips} ZIPs, {total_episodes} episodes")
    logger.info(f"Output: {output_dir}")


if __name__ == "__main__":
    main()