import json
from pathlib import Path

import h5py
import numpy as np
from PIL import Image


def get_task_info(task_json_path: str | Path) -> list[dict]:
    with open(task_json_path, "r") as f:
        task_info: list = json.load(f)
    task_info.sort(key=lambda episode: episode["episode_id"])
    return task_info


def get_task_id(task_json_path: str | Path) -> str:
    return Path(task_json_path).stem.split("_")[-1]


def get_episode_ids(src_path: str | Path, task_id: str | int) -> list[int]:
    observations_dir = Path(src_path) / "observations" / str(task_id)
    return sorted(
        int(path.name)
        for path in observations_dir.glob("*")
        if path.is_dir() and path.name.isdigit()
    )


def get_episode_videos(
    src_path: str | Path,
    task_id: str | int,
    episode_id: int,
    agibot_world_config: dict,
) -> dict[str, Path]:
    ob_dir = Path(src_path) / f"observations/{task_id}/{episode_id}"
    return {
        f"observation.images.{key}": ob_dir / "videos" / f"{key}_color.mp4"
        if "sensor" not in key
        else ob_dir / "tactile" / f"{key}.mp4"  # HACK: handle tactile videos
        for key in agibot_world_config["images"]
        if "depth" not in key
    }


def has_episode_videos(
    src_path: str | Path,
    task_id: str | int,
    episode_id: int,
    agibot_world_config: dict,
) -> bool:
    videos = get_episode_videos(src_path, task_id, episode_id, agibot_world_config)
    return all(video_path.exists() for video_path in videos.values())


def load_depths(root_dir: str, camera_name: str):
    cam_path = Path(root_dir)
    all_imgs = sorted(list(cam_path.glob(f"{camera_name}*")))
    return [
        np.array(Image.open(f)).astype(np.float32)[:, :, None] / 1000 for f in all_imgs
    ]


def load_local_dataset(
    episode_id: int,
    src_path: str | Path,
    task_id: str | int,
    save_depth: bool,
    AgiBotWorld_CONFIG: dict,
) -> tuple[int, list[dict], dict[str, Path]] | None:
    """Load local dataset and return a dict with observations and actions"""
    ob_dir = Path(src_path) / f"observations/{task_id}/{episode_id}"
    proprio_dir = Path(src_path) / f"proprio_stats/{task_id}/{episode_id}"

    state = {}
    action = {}
    with h5py.File(proprio_dir / "proprio_stats.h5", "r") as f:
        for key in AgiBotWorld_CONFIG["states"]:
            state[f"observation.states.{key}"] = np.array(
                f["state/" + key.replace(".", "/")], dtype=np.float32
            )
        for key in AgiBotWorld_CONFIG["actions"]:
            action[f"actions.{key}"] = np.array(
                f["action/" + key.replace(".", "/")], dtype=np.float32
            )

        # HACK: agibot team forgot to pad or filter some of the values
        num_frames = len(next(iter(state.values())))
        for action_key, action_value in action.items():
            if 0 == len(action_value):
                print("0 action occurs, padding all with zeros later")
            elif len(action_value) < num_frames:
                state_key = action_key.replace("actions", "state").replace(".", "/")
                new_action_value = np.array(f[state_key], dtype=np.float32).copy()
                action_index_key = "/".join(
                    list(action_key.replace("actions", "action").split(".")[:-1])
                    + ["index"]
                )
                action_index = np.array(f[action_index_key])
                # agibot lost end index, replace it with joint
                if not action_index.size:
                    action_index_key = action_index_key.replace("end", "joint")
                    action_index = np.array(f[action_index_key])
                new_action_value[action_index] = action_value
                action[action_key] = new_action_value
            elif len(action_value) > num_frames:
                print("corrupt data, skipping")
                return None

    if save_depth:
        depth_imgs = load_depths(ob_dir / "depth", "head_depth")
        assert num_frames == len(depth_imgs), (
            "Number of images and states are not equal"
        )

    state_key_prefix_len = len("observation.states.")
    action_key_prefix_len = len("actions.")
    frames = [
        {
            **({"observation.images.head_depth": depth_imgs[i]} if save_depth else {}),
            **{
                key: value[i]
                if value.size
                else np.zeros(
                    AgiBotWorld_CONFIG["states"][key[state_key_prefix_len:]]["shape"],
                    dtype=AgiBotWorld_CONFIG["states"][key[state_key_prefix_len:]][
                        "dtype"
                    ],
                )
                for key, value in state.items()
            },
            **{
                key: value[i]
                if value.size
                else np.zeros(
                    AgiBotWorld_CONFIG["actions"][key[action_key_prefix_len:]]["shape"],
                    dtype=AgiBotWorld_CONFIG["actions"][key[action_key_prefix_len:]][
                        "dtype"
                    ],
                )
                for key, value in action.items()
            },
        }
        for i in range(num_frames)
    ]

    videos = get_episode_videos(src_path, task_id, episode_id, AgiBotWorld_CONFIG)
    return episode_id, frames, videos