import os
import re
import h5py
import json
import glob
import numpy as np
from pathlib import Path
import argparse
import random

def get_task_name_from_path(path):
    match = re.search(r'/data/([^/]+)/', path)
    if match:
        return match.group(1)
    else:
        return "unknown_task"

def generate_llm_items(task_name, task_config, episode_id, instruction_dict, robot_obs, future_step_num=5, image_root_path=None):
    episode_str = f"{task_name}_episode{episode_id:02d}"
    instructions = instruction_dict.get("seen", ["Do the task."])

    llm_items = []
    num_steps = robot_obs.shape[0]

    for t in range(num_steps):
        instruction = random.choice(instructions)

        current_robot_obs = robot_obs[t]
        robot_obs_string = " ".join(map(str, current_robot_obs.flatten()))

        # Get future five steps of robot_obs data (14×5=70 dimensions)
        future_indices = list(range(t + 1, min(t + future_step_num + 1, num_steps)))
        future_robot_obs = []
        for idx in future_indices:
            future_robot_obs.append(robot_obs[idx])
        # Pad with last frame if fewer than five steps remain
        while len(future_robot_obs) < future_step_num:
            future_robot_obs.append(robot_obs[-1])
        # Combine into 70-dimensional vector
        future_robot_obs = np.array(future_robot_obs).flatten()
        future_robot_obs_string = " ".join(map(str, future_robot_obs))

        # Use new image path: episode{episode_id}/{frame_idx:02d}.jpg
        frame_idx = t + 1
        image_path = None
        if image_root_path is not None:
            image_path = str(Path(task_name) / task_config / f"episode{episode_id}" / f"{frame_idx:02d}.jpg")

        llm_item = {
            "id": f"{episode_str}_{t:03d}",
            "image": image_path,  # Single combined image path
            "conversations": [
                {
                    "from": "human",
                    "value": "<image>\n" + instruction + "\n" + robot_obs_string
                },
                {
                    "from": "gpt",
                    "value": future_robot_obs_string  # 70-dimensional future five steps
                }
            ],
            "embody": True
        }

        llm_items.append(llm_item)

    return llm_items

def generate_llm_dataset(data_root_dir, instruction_dir, image_root_path, save_root_dir, task_name=None, task_config=None,future_step_num=5):
    if task_name is None:
        task_name = get_task_name_from_path(data_root_dir)
    print(f"Task name: {task_name}")

    episode_files = sorted(
        glob.glob(os.path.join(data_root_dir, "episode*.hdf5")),
        key=lambda x: int(re.search(r"episode(\d+)\.hdf5", os.path.basename(x)).group(1)) if re.search(r"episode(\d+)\.hdf5", os.path.basename(x)) else float('inf')
    )

    if len(episode_files) == 0:
        print(f"No episode*.hdf5 files found for task {task_name}")
        return

    save_dir = os.path.join(save_root_dir, task_name, task_config)
    os.makedirs(save_dir, exist_ok=True)

    for file_path in episode_files:
        match = re.search(r"episode(\d+)\.hdf5", os.path.basename(file_path))
        if not match:
            print(f"Unrecognized filename format, skipping: {os.path.basename(file_path)}")
            continue

        episode_id = int(match.group(1))
        
        instruction_path = os.path.join(instruction_dir, f"episode{episode_id}.json")
        if not os.path.exists(instruction_path):
            print(f"Missing instruction file for episode {episode_id}, skipping")
            continue

        try:
            with open(instruction_path, "r") as f:
                instruction_dict = json.load(f)

            with h5py.File(file_path, "r") as h5f:
                # Read only joint_action/vector
                robot_obs = h5f["joint_action/vector"][()]  # shape (steps, 14)

            llm_items = generate_llm_items(
                task_name=task_name,
                task_config=task_config,
                episode_id=episode_id,
                instruction_dict=instruction_dict,
                robot_obs=robot_obs,
                future_step_num=future_step_num,
                image_root_path=image_root_path
            )

            # Save each episode as a JSON file
            save_file = os.path.join(save_dir, f"episode{episode_id}.json")
            with open(save_file, "w", encoding="utf-8") as f_out:
                json.dump(llm_items, f_out, ensure_ascii=False, indent=2)

            print(f"Episode {episode_id} processed successfully!")

        except Exception as e:
            print(f"Failed to process episode {episode_id}: {e}")

    print(f"All episodes for task {task_name} processed successfully!")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate LLM-format dataset from HDF5 and images")
    parser.add_argument("--task_name", type=str, default=None, help="Task name")
    parser.add_argument("--task_config", type=str, default=None, help="Task config")
    parser.add_argument("--data_root_dir", type=str, required=True, help="HDF5 data directory")
    parser.add_argument("--instruction_dir", type=str, required=True, help="Instruction JSON directory")
    parser.add_argument("--image_root_path", type=str, required=True, help="Image root path (absolute path)")
    parser.add_argument("--save_root_dir", type=str, required=True, help="Root directory for saving generated JSON files")
    parser.add_argument("--future_step_num", type=int, default=5, help="Number of future frame steps")
    args = parser.parse_args()

    generate_llm_dataset(
        data_root_dir=args.data_root_dir,
        instruction_dir=args.instruction_dir,
        image_root_path=args.image_root_path,
        save_root_dir=args.save_root_dir,
        task_name=args.task_name,
        task_config=args.task_config,
        future_step_num=args.future_step_num
    )
