import os
import re
import h5py
import numpy as np
import cv2
import glob
import argparse
from PIL import Image

def get_task_name(path):
    """
    Extract task name from path, assuming /data/<task_name>/ format
    """
    match = re.search(r'/data/([^/]+)/', path)
    if match:
        return match.group(1)
    else:
        return "unknown_task"
    

def extract_images(hdf5_path, save_root_dir, task_name, task_config, episode_id=0):
    """
    Directory structure:
    save_root_dir/task_name/episode{episode_id}/{01..89}.jpg
    Each frame combines four camera images (left_camera, right_camera, head_camera, front_camera) into a 2x2 grid in RGB format

    :param hdf5_path: Path to HDF5 file
    :param save_root_dir: Root path for saving images
    :param task_name: Task name string
    :param episode_id: Episode number, e.g., 0
    """
    cameras = ["left_camera", "right_camera", "head_camera", "front_camera"]
    episode_folder = os.path.join(save_root_dir, task_name, task_config, f"episode{episode_id}")
    os.makedirs(episode_folder, exist_ok=True)

    with h5py.File(hdf5_path, "r") as f:
        num_frames = len(f[f"observation/{cameras[0]}/rgb"])  # Use left_camera frame count as reference
        for i in range(num_frames):
            # Image path: episode{episode_id}/{frame_idx:02d}.jpg
            frame_idx = i + 1
            img_path = os.path.join(episode_folder, f"{frame_idx:02d}.jpg")

            # Load and decode images from all cameras
            images = []
            for cam in cameras:
                img_bytes = f[f"observation/{cam}/rgb"][i]
                img_array = np.frombuffer(img_bytes, dtype=np.uint8)
                img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

                if img is None:
                    print(f"[Warning] episode{episode_id} frame {frame_idx} {cam} image decoding failed, skipping this frame")
                    break  # Skip frame if any image fails to decode
                # Convert BGR to RGB immediately after decoding
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                images.append(img)
            else:
                # Combine images only if all are successfully loaded
                if len(images) == len(cameras):
                    # Get image dimensions (assuming all images have the same size)
                    height, width = images[0].shape[:2]

                    # Create a 2x2 grid canvas (RGB)
                    combined_img = np.zeros((height * 2, width * 2, 3), dtype=np.uint8)

                    # Place images in 2x2 grid
                    # Top-left: left_camera (index 0)
                    combined_img[0:height, 0:width] = images[0]
                    # Top-right: right_camera (index 1)
                    combined_img[0:height, width:width*2] = images[1]
                    # Bottom-left: head_camera (index 2)
                    combined_img[height:height*2, 0:width] = images[2]
                    # Bottom-right: front_camera (index 3)
                    combined_img[height:height*2, width:width*2] = images[3]

                    # Save combined image as RGB using PIL
                    combined_img_pil = Image.fromarray(combined_img)
                    combined_img_pil.save(img_path, "JPEG")
                else:
                    print(f"[Warning] episode{episode_id} frame {frame_idx} missing some images, skipping combination")

    print(f" episode{episode_id} processed successfuly")

def process_episode_files(data_root_dir, save_root_dir, task_config, task_name=None):
    """
    Batch process all episode*.hdf5 files in a directory.

    :param data_root_dir: Directory of raw HDF5 files, e.g., /.../grab_roller/clean/data/
    :param save_root_dir: Root directory for saving images
    :param task_name: Task name, if None, extract from path
    """
    if task_name is None:
        task_name = get_task_name(data_root_dir)
    print(f"Task name: {task_name}")

    # Get episode files and sort by numerical order
    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 len(episode_files) == 0:
        print("No episode*.hdf5 files found")
        return

    print(f"Found {len(episode_files)} episode files, starting processing...")

    for file_path in episode_files:
        match = re.search(r"episode(\d+)\.hdf5", os.path.basename(file_path))
        if match:
            episode_id = int(match.group(1))
            try:
                extract_images(
                    hdf5_path=file_path,
                    save_root_dir=save_root_dir,
                    task_name=task_name,
                    task_config=task_config,
                    episode_id=episode_id
                )
            except Exception as e:
                print(f" Failed to process {file_path}: {e}")
        else:
            print(f" Unrecognized filename format: {file_path}")

    print(" All episode processing completed!")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process HDF5 episodes and extract 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("--save_root_dir", type=str, required=True, help="Root directory for saving images")
    args = parser.parse_args()

    process_episode_files(args.data_root_dir, args.save_root_dir, args.task_config, args.task_name)