from __future__ import annotations

from dataclasses import dataclass
from time import perf_counter, sleep
from typing import Iterable, Optional

import mujoco
import numpy as np
import pinocchio as pin

from .kinematics import validate_se3
from .mujoco_model import build_normalized_dual_mjcf
from .types import JointLimits, physical_joint_limits


CONTROLLED_HOME_Q_RAD = np.deg2rad([0.0, 30.0, 0.0, 60.0, 0.0, 60.0, 0.0])
PROJECT_INITIAL_Q_RAD = {
    "left": np.deg2rad([-167.21, 28.48, 28.21, 61.35, -14.40, 84.49, -124.51]),
    "right": np.deg2rad([-25.60, 34.09, -19.55, 71.59, 16.97, 80.98, 59.67]),
}


@dataclass(frozen=True)
class PlaybackResult:
    samples: int
    elapsed_sec: float
    max_joint_step_rad: float
    final_flange_pose: pin.SE3


class DualArmMuJoCo:
    """Threadless, kinematic MuJoCo backend for the normalized dual RM75 scene."""

    def __init__(
        self,
        controlled_arm: str = "left",
        inactive_q_rad: Optional[np.ndarray] = None,
        limits: Optional[JointLimits] = None,
    ) -> None:
        if controlled_arm not in {"left", "right"}:
            raise ValueError("controlled_arm must be 'left' or 'right'")
        self.controlled_arm = controlled_arm
        self.inactive_arm = "right" if controlled_arm == "left" else "left"
        self.limits = limits or physical_joint_limits()
        self.mjcf_xml, self.assets = build_normalized_dual_mjcf()
        self.model = mujoco.MjModel.from_xml_string(self.mjcf_xml, self.assets)
        self.data = mujoco.MjData(self.model)
        self._qpos_addresses = {
            arm: np.array(
                [
                    self.model.jnt_qposadr[
                        mujoco.mj_name2id(
                            self.model,
                            mujoco.mjtObj.mjOBJ_JOINT,
                            f"{arm}_joint_{joint_index}",
                        )
                    ]
                    for joint_index in range(1, 8)
                ],
                dtype=int,
            )
            for arm in ("left", "right")
        }
        self._dof_addresses = {
            arm: np.array(
                [
                    self.model.jnt_dofadr[
                        mujoco.mj_name2id(
                            self.model,
                            mujoco.mjtObj.mjOBJ_JOINT,
                            f"{arm}_joint_{joint_index}",
                        )
                    ]
                    for joint_index in range(1, 8)
                ],
                dtype=int,
            )
            for arm in ("left", "right")
        }
        self._site_ids = {
            arm: mujoco.mj_name2id(
                self.model, mujoco.mjtObj.mjOBJ_SITE, f"{arm}_flange"
            )
            for arm in ("left", "right")
        }
        self._target_mocap_ids = {}
        for arm in ("left", "right"):
            marker_body_id = mujoco.mj_name2id(
                self.model, mujoco.mjtObj.mjOBJ_BODY, f"{arm}_target_marker"
            )
            mocap_id = int(self.model.body_mocapid[marker_body_id])
            if mocap_id < 0:
                raise ValueError(f"{arm}_target_marker is not a MuJoCo mocap body")
            self._target_mocap_ids[arm] = mocap_id

        inactive = (
            PROJECT_INITIAL_Q_RAD[self.inactive_arm]
            if inactive_q_rad is None
            else np.asarray(inactive_q_rad, dtype=float)
        )
        self.set_arm_configuration(self.inactive_arm, inactive)
        self.set_arm_configuration(self.controlled_arm, CONTROLLED_HOME_Q_RAD)
        for arm in ("left", "right"):
            self.set_arm_target_marker(arm, self.get_flange_pose(arm))
        self._manual_inactive_q: Optional[np.ndarray] = None

    @staticmethod
    def _validate_arm(arm: str) -> None:
        if arm not in {"left", "right"}:
            raise ValueError("arm must be 'left' or 'right'")

    def _validate_q(self, q_rad: np.ndarray) -> np.ndarray:
        q = np.asarray(q_rad, dtype=float)
        if q.shape != (7,):
            raise ValueError(f"arm configuration must have shape (7,), got {q.shape}")
        if not np.all(np.isfinite(q)):
            raise ValueError("arm configuration must be finite")
        if not self.limits.contains(q):
            raise ValueError(f"configuration is outside {self.limits.name} joint limits")
        return q.copy()

    def set_arm_configuration(self, arm: str, q_rad: np.ndarray) -> None:
        self._validate_arm(arm)
        q = self._validate_q(q_rad)
        self.data.qpos[self._qpos_addresses[arm]] = q
        self.data.qvel[:] = 0.0
        mujoco.mj_forward(self.model, self.data)

    def set_dual_configuration(
        self, left_q_rad: np.ndarray, right_q_rad: np.ndarray
    ) -> None:
        left_q = self._validate_q(left_q_rad)
        right_q = self._validate_q(right_q_rad)
        self.data.qpos[self._qpos_addresses["left"]] = left_q
        self.data.qpos[self._qpos_addresses["right"]] = right_q
        self.data.qvel[:] = 0.0
        mujoco.mj_forward(self.model, self.data)

    def get_arm_configuration(self, arm: str) -> np.ndarray:
        self._validate_arm(arm)
        return self.data.qpos[self._qpos_addresses[arm]].copy()

    def get_flange_pose(self, arm: str) -> pin.SE3:
        self._validate_arm(arm)
        site_id = self._site_ids[arm]
        return pin.SE3(
            self.data.site_xmat[site_id].reshape(3, 3).copy(),
            self.data.site_xpos[site_id].copy(),
        )

    def set_target_marker(self, target_se3: pin.SE3) -> None:
        self.set_arm_target_marker(self.controlled_arm, target_se3)

    def set_arm_target_marker(self, arm: str, target_se3: pin.SE3) -> None:
        self._validate_arm(arm)
        validate_se3(target_se3, "target_se3")
        quaternion = pin.Quaternion(target_se3.rotation)
        mocap_id = self._target_mocap_ids[arm]
        self.data.mocap_pos[mocap_id] = target_se3.translation
        self.data.mocap_quat[mocap_id] = [
            quaternion.w,
            quaternion.x,
            quaternion.y,
            quaternion.z,
        ]
        mujoco.mj_forward(self.model, self.data)

    def _validated_trajectory(self, q_trajectory: Iterable[np.ndarray]) -> np.ndarray:
        trajectory = np.asarray(list(q_trajectory), dtype=float)
        if trajectory.ndim != 2 or trajectory.shape[1:] != (7,):
            raise ValueError("q_trajectory must have shape (N, 7)")
        if trajectory.shape[0] == 0:
            raise ValueError("q_trajectory must contain at least one sample")
        if not np.all(np.isfinite(trajectory)):
            raise ValueError("q_trajectory must be finite")
        if np.any(trajectory < self.limits.lower) or np.any(
            trajectory > self.limits.upper
        ):
            raise ValueError(f"trajectory exceeds {self.limits.name} joint limits")
        return trajectory

    def play_trajectory(
        self,
        q_trajectory: Iterable[np.ndarray],
        dt: float = 1.0 / 90.0,
        realtime: bool = False,
        viewer=None,
    ) -> PlaybackResult:
        if not np.isfinite(dt) or dt <= 0.0:
            raise ValueError("dt must be finite and positive")
        trajectory = self._validated_trajectory(q_trajectory)
        started = perf_counter()
        previous = self.get_arm_configuration(self.controlled_arm)
        max_step = 0.0
        for q in trajectory:
            frame_started = perf_counter()
            max_step = max(max_step, float(np.max(np.abs(q - previous))))
            self.set_arm_configuration(self.controlled_arm, q)
            previous = q
            if viewer is not None:
                viewer.sync()
            if realtime:
                remaining = dt - (perf_counter() - frame_started)
                if remaining > 0.0:
                    sleep(remaining)
        return PlaybackResult(
            samples=len(trajectory),
            elapsed_sec=perf_counter() - started,
            max_joint_step_rad=max_step,
            final_flange_pose=self.get_flange_pose(self.controlled_arm),
        )

    def configure_manual_drag(self, damping: float = 1.5) -> None:
        if not np.isfinite(damping) or damping <= 0.0:
            raise ValueError("manual-drag damping must be finite and positive")
        self.model.opt.gravity[:] = 0.0
        self.model.dof_damping[self._dof_addresses[self.controlled_arm]] = damping
        self.data.qvel[:] = 0.0
        self.data.xfrc_applied[:] = 0.0
        self._manual_inactive_q = self.get_arm_configuration(self.inactive_arm)
        mujoco.mj_forward(self.model, self.data)

    def step_manual_drag(self) -> None:
        if self._manual_inactive_q is None:
            raise RuntimeError("configure_manual_drag() must be called first")
        mujoco.mj_step(self.model, self.data)
        self.data.qpos[self._qpos_addresses[self.inactive_arm]] = self._manual_inactive_q
        self.data.qvel[self._dof_addresses[self.inactive_arm]] = 0.0
        mujoco.mj_forward(self.model, self.data)

    def render(
        self,
        width: int = 1280,
        height: int = 720,
        *,
        segmentation: bool = False,
    ) -> np.ndarray:
        if width <= 0 or height <= 0:
            raise ValueError("render dimensions must be positive")
        renderer = mujoco.Renderer(self.model, height=height, width=width)
        try:
            if segmentation:
                renderer.enable_segmentation_rendering()
            renderer.update_scene(self.data, camera="overview")
            return renderer.render().copy()
        finally:
            renderer.close()
