"""Env-adapter IO helpers for single-arm LIBERO tasks.

Env code extracts a compact ndarray payload from raw LIBERO observations.
Adapters then build proprioception, masks, and 7-dof right-arm action chunks
from that payload. This module intentionally depends only on NumPy so adapter
processes can import it without importing robosuite or LIBERO.
"""

from __future__ import annotations

import math
import os

import numpy as np


# Fallback single-arm LIBERO dof_config when train_config does not provide one.
_LIBERO_FALLBACK_DOF_CONFIG = {
    "follow_right_ee_cartesian_pos": 3,
    "follow_right_ee_rotation": 3,
    "follow_right_gripper": 1,
}

_LIBERO_FALLBACK_AGENT_POS_CONFIG = dict(_LIBERO_FALLBACK_DOF_CONFIG)
_VIRTUAL_TAIL_KEYS = frozenset(("action_padding",))


def _resolve_dof_config(train_config: dict) -> dict:
    return (
        train_config.get("dof_config")
        or train_config.get("task", {}).get("dof_config")
        or _LIBERO_FALLBACK_DOF_CONFIG
    )


def _resolve_agent_pos_config(train_config: dict) -> dict:
    return (
        train_config.get("agent_pos_config")
        or train_config.get("task", {}).get("agent_pos_config")
        or _LIBERO_FALLBACK_AGENT_POS_CONFIG
    )


def _move_virtual_keys_to_tail(layout: dict) -> dict:
    head = {k: v for k, v in layout.items() if k not in _VIRTUAL_TAIL_KEYS}
    tail = {k: v for k, v in layout.items() if k in _VIRTUAL_TAIL_KEYS}
    return {**head, **tail}


def _effective_agent_pos_config(train_config: dict, state_values: dict) -> dict:
    config = _move_virtual_keys_to_tail(dict(_resolve_agent_pos_config(train_config)))
    gripper_key = next(
        (
            key
            for key in config
            if key.replace("follow_", "").replace("master_", "") == "right_gripper"
        ),
        None,
    )
    if gripper_key is None:
        return config

    old_dim = int(config[gripper_key])
    target_dim = None
    override = os.environ.get("WALLX_LIBERO_STATE_GRIPPER_DIM")
    if override:
        target_dim = int(override)
    elif os.environ.get("WALLX_LIBERO_AUTO_STATE_GRIPPER_DIM", "1") != "0":
        norm_dim = int(train_config.get("_libero_proprio_norm_dim") or 0)
        real_dim = sum(v for k, v in config.items() if k not in _VIRTUAL_TAIL_KEYS)
        if norm_dim == real_dim + 1:
            target_dim = old_dim + 1

    available_dim = state_values["right_gripper"].shape[1]
    if target_dim is None or target_dim == old_dim or target_dim > available_dim:
        return config

    config[gripper_key] = target_dim
    delta = target_dim - old_dim
    if "action_padding" in config:
        config["action_padding"] = max(0, int(config["action_padding"]) - delta)
    return config


def _build_right_arm_state_values(obs_ndarrays: dict) -> dict[str, np.ndarray]:
    """Bare-key state tensors for single-arm LIBERO proprio construction."""
    rot3 = np.asarray(obs_ndarrays["eef_axisangle"], dtype=np.float32).reshape(1, 3)
    values: dict[str, np.ndarray] = {
        "right_ee_cartesian_pos": np.asarray(
            obs_ndarrays["eef_pos"], dtype=np.float32
        ).reshape(1, 3),
        "right_ee_rotation": rot3,
        "right_gripper": np.asarray(obs_ndarrays["gripper"], dtype=np.float32).reshape(
            1, -1
        ),
    }
    from wall_x._vendor.x2robot_utils.geometry import euler_to_matrix_zyx_6d_nb

    rot6d = euler_to_matrix_zyx_6d_nb(rot3.astype(np.float64)).reshape(1, 6)
    values["right_ee_rotation_6D"] = rot6d.astype(np.float32)
    return values


# Auxiliary action keys that should be masked out for single-arm LIBERO.
_DOF_MASK_ZERO_KEYS = frozenset(
    (
        "follow_left_ee_cartesian_pos",
        "follow_left_ee_rotation",
        "follow_left_ee_rotation_6D",
        "follow_left_gripper",
        "head_actions",
        "height",
        "velocity_decomposed",
        "action_padding",
    )
)


# ============================================================
# LIBERO raw observation decoding helpers.
# ============================================================


def _get_libero_image(obs: dict) -> np.ndarray:
    """Return the third-person camera image, rotated to match preprocessing."""
    return obs["agentview_image"][::-1, ::-1]


def get_rollout_frame(obs: dict) -> np.ndarray:
    """Return one RGB frame for rollout MP4 saving."""
    return np.asarray(_get_libero_image(obs), dtype=np.uint8)


def _get_libero_wrist_image(obs: dict) -> np.ndarray:
    """Return the wrist camera image, rotated to match preprocessing."""
    return obs["robot0_eye_in_hand_image"][::-1, ::-1]


def _quat2axisangle(quat) -> np.ndarray:
    """Convert an xyzw quaternion to a 3D axis-angle vector."""
    if quat[3] > 1.0:
        quat[3] = 1.0
    elif quat[3] < -1.0:
        quat[3] = -1.0

    den = np.sqrt(1.0 - quat[3] * quat[3])
    if math.isclose(den, 0.0):
        return np.zeros(3)

    return (quat[:3] * 2.0 * math.acos(quat[3])) / den


# ============================================================
# Env-side to adapter-side observation encoding.
# ============================================================


def encode_raw_obs(raw_obs: dict) -> dict:
    """Extract the minimal ndarray payload from a raw LIBERO observation.

    The payload contains three 1-D state arrays and two rotated image arrays.
    """
    if "agentview_image" not in raw_obs:
        raise KeyError(
            "agentview_image missing in raw_obs; render-skip may have returned "
            "a stale observation"
        )
    return {
        "eef_pos": np.asarray(raw_obs["robot0_eef_pos"], dtype=np.float32),
        "eef_axisangle": np.asarray(
            _quat2axisangle(raw_obs["robot0_eef_quat"]), dtype=np.float32
        ),
        "gripper": np.asarray(raw_obs["robot0_gripper_qpos"], dtype=np.float32),
        "face_view": _get_libero_image(raw_obs),
        "wrist_view": _get_libero_wrist_image(raw_obs),
    }


# ============================================================
# Adapter-side proprioception and mask construction.
# ============================================================


def encode_proprio(
    obs_ndarrays: dict,
    train_config: dict,
    action_horizon: int,
) -> dict:
    """Convert a single-arm LIBERO ndarray payload into model input fields.

    Returned fields include proprioception, agent_pos_mask, dof_mask,
    face_view, and right_wrist_view.
    """
    state_values = _build_right_arm_state_values(obs_ndarrays)
    agent_pos_config = _effective_agent_pos_config(train_config, state_values)
    dof_config = _move_virtual_keys_to_tail(dict(_resolve_dof_config(train_config)))

    propri_parts: list[np.ndarray] = []
    mask_parts: list[np.ndarray] = []
    for key, dim in agent_pos_config.items():
        bare = key.replace("follow_", "").replace("master_", "")
        if bare in state_values:
            v = state_values[bare]
            if bare == "right_gripper" and v.shape[1] > dim:
                v = v[:, :dim]
            if v.shape[1] != dim:
                raise ValueError(
                    f"agent_pos_config[{key!r}]={dim} does not match "
                    f"observation dimension {v.shape[1]}"
                )
            propri_parts.append(v)
            mask_parts.append(np.ones((1, dim), dtype=np.float32))
        else:
            propri_parts.append(np.zeros((1, dim), dtype=np.float32))
            mask_parts.append(np.zeros((1, dim), dtype=np.float32))

    # (1, 1, D)
    proprioception = np.concatenate(propri_parts, axis=1)[None]
    agent_pos_mask = np.concatenate(mask_parts, axis=1)[None]

    # dof_mask: (1, T, D_action)
    total_dof = sum(dof_config.values())
    dof_mask = np.ones((1, action_horizon, total_dof))
    start = 0
    for key, dim in dof_config.items():
        if key in _DOF_MASK_ZERO_KEYS:
            dof_mask[:, :, start : start + dim] = 0
        start += dim

    return {
        "proprioception": proprioception.astype(np.float32),
        "agent_pos_mask": agent_pos_mask.astype(np.float32),
        "dof_mask": dof_mask,
        "face_view": obs_ndarrays["face_view"],
        "right_wrist_view": obs_ndarrays["wrist_view"],
    }


# ============================================================
# Adapter-side action decoding.
# ============================================================


def decode_chunk(predict_action: np.ndarray, train_config: dict) -> np.ndarray:
    """Extract a 7-dof right-arm chunk from model action output."""
    if predict_action.ndim == 3:
        predict_action = predict_action[0]

    dof_config = _resolve_dof_config(train_config)
    slices: dict[str, slice] = {}
    start = 0
    for key, dim in dof_config.items():
        bare = key.replace("follow_", "").replace("master_", "")
        slices[bare] = slice(start, start + dim)
        start += dim

    pos = predict_action[:, slices["right_ee_cartesian_pos"]]
    grip = predict_action[:, slices["right_gripper"]]

    if "right_ee_rotation_6D" in slices:
        from wall_x._vendor.x2robot_utils.geometry import so3_to_euler_zyx_batch_nb

        rot6d = np.asarray(
            predict_action[:, slices["right_ee_rotation_6D"]], dtype=np.float64
        )
        rot = so3_to_euler_zyx_batch_nb(rot6d).astype(np.float32)
    elif "right_ee_rotation" in slices:
        rot = predict_action[:, slices["right_ee_rotation"]]
    else:
        raise KeyError(
            "dof_config has no right-arm rotation slice "
            f"(keys={list(slices.keys())})"
        )

    return np.concatenate([pos, rot, grip], axis=1)


def gripper_model_to_libero_osc(grip_2d: np.ndarray) -> float:
    """Map model gripper output to robosuite OSC_POSE gripper command in [-1, 1]."""
    g = np.asarray(grip_2d, dtype=np.float64).reshape(-1)
    if g.size == 0:
        raise ValueError("empty gripper action")
    cmd = float(g[0])
    if os.environ.get("WALLX_LIBERO_GRIPPER_BINARIZE", "0") == "1":
        if abs(cmd) < 1e-6:
            cmd = -1.0
        else:
            cmd = float(np.sign(cmd))
    if os.environ.get("WALLX_LIBERO_INVERT_GRIPPER", "0") == "1":
        cmd *= -1.0
    return cmd


def _sanitize_task_description(task_description: str, max_len: int = 50) -> str:
    return (
        task_description.lower()
        .replace(" ", "_")
        .replace("\n", "_")
        .replace(".", "_")[:max_len]
    )


def save_rollout_video(
    rollout_dir: str,
    rollout_images: list[np.ndarray],
    *,
    task_id: int,
    episode_idx: int,
    success: bool,
    task_description: str,
    fps: int = 30,
) -> str | None:
    """Save an MP4 replay of one LIBERO episode."""
    if not rollout_images:
        return None

    import imageio

    os.makedirs(rollout_dir, exist_ok=True)
    task_slug = _sanitize_task_description(task_description)
    mp4_path = os.path.join(
        rollout_dir,
        f"task{task_id}_ep{episode_idx}--success={int(success)}--{task_slug}.mp4",
    )
    writer = imageio.get_writer(mp4_path, fps=fps, macro_block_size=1)
    try:
        for img in rollout_images:
            writer.append_data(np.asarray(img, dtype=np.uint8))
    finally:
        writer.close()
    return mp4_path


def model_action_to_libero_env(action: np.ndarray) -> np.ndarray:
    """Convert model output to the robosuite OSC_POSE 7D action.

    Internal LIBERO evaluation passes the model's 7D chunk directly to
    ``env.step``. Keep that as the public default; conversion modes remain
    available only for ablations through environment variables.
    """
    from scipy.spatial.transform import Rotation as R

    a = np.asarray(action, dtype=np.float64).reshape(-1)
    if a.size not in (7, 8):
        raise ValueError(f"expected 7D or 8D model action, got shape {a.shape}")

    pos_delta = a[:3]
    rot_mode = os.environ.get("WALLX_LIBERO_ROT_MODE", "direct").strip()
    if rot_mode == "euler_zyx_to_rotvec":
        rot_aa = R.from_euler("zyx", a[3:6]).as_rotvec()
    elif rot_mode == "direct":
        rot_aa = a[3:6]
    else:
        raise ValueError(
            "WALLX_LIBERO_ROT_MODE must be 'euler_zyx_to_rotvec' or 'direct', "
            f"got {rot_mode!r}"
        )
    grip = gripper_model_to_libero_osc(a[6:8] if a.size == 8 else a[6:7])
    return np.concatenate([pos_delta, rot_aa, np.array([grip], dtype=np.float64)])