from __future__ import annotations

from contextlib import nullcontext
from dataclasses import dataclass
from time import sleep
from typing import Callable, Collection, Dict, Mapping, Optional

import numpy as np
import pinocchio as pin

from .dual_arm import DualArmAssembly
from .kinematics import RM75Kinematics, pose_errors, validate_se3
from .mujoco_backend import DualArmMuJoCo
from .robot_backend import DualArmJointState
from .solver import RM75IkSolver, deterministic_recovery_seeds
from .teleop_config import ArmName, ArmTeleopProfile
from .types import IkOptions, physical_joint_limits, teleop_joint_limits


@dataclass(frozen=True)
class InitialPoseDiagnostic:
    arm: ArmName
    solved_q_rad: np.ndarray
    solved_position_error_m: float
    solved_orientation_error_rad: float
    configured_joint_position_error_m: float
    configured_joint_orientation_error_rad: float


class MujocoRobot:
    """Dual-arm kinematic backend implementing the backend-neutral robot contract."""

    def __init__(
        self,
        profiles: Mapping[ArmName, ArmTeleopProfile],
        *,
        initialize_from_tcp: bool = True,
    ) -> None:
        if set(profiles) != {"left", "right"}:
            raise ValueError("profiles must contain left and right arms")
        self.profiles = dict(profiles)
        self.scene = DualArmMuJoCo(controlled_arm="left", limits=teleop_joint_limits())
        self.assembly = DualArmAssembly.from_source_urdf()
        self._kinematics = {
            arm: RM75Kinematics(limits=teleop_joint_limits())
            for arm in ("left", "right")
        }
        self._viewer = None
        self._closed = False
        self._stopped_arms: set[ArmName] = set()
        self._initial_joint_positions: Optional[Dict[ArmName, np.ndarray]] = None
        self.initial_pose_diagnostics: Dict[ArmName, InitialPoseDiagnostic] = {}
        if initialize_from_tcp:
            self._initialize_from_configured_tcp()

    def _initialize_from_configured_tcp(self) -> None:
        solved: Dict[ArmName, np.ndarray] = {}
        limits = teleop_joint_limits()
        physical_kinematics = RM75Kinematics(limits=physical_joint_limits())
        for seed_offset, arm in enumerate(("left", "right")):
            profile = self.profiles[arm]
            kinematics = RM75Kinematics(limits=limits)
            solver = RM75IkSolver(kinematics)
            flange_target = profile.initial_tcp * profile.tool_from_flange.inverse()
            seeds = deterministic_recovery_seeds(
                limits, count=24, random_seed=750 + seed_offset
            )
            if limits.contains(profile.configured_initial_q_rad):
                seeds.insert(0, profile.configured_initial_q_rad)
            result = solver.solve_multistart(
                flange_target,
                seeds,
                IkOptions(max_iterations=700, time_limit_sec=None),
            )
            if not result.success or result.q is None:
                raise RuntimeError(
                    f"failed to resolve {arm} initial_tcp_pose: "
                    f"{result.status.value}: {result.message}"
                )
            solved[arm] = result.q.copy()
            solved_tcp = kinematics.forward(result.q, profile.tool_from_flange)
            solved_error = pose_errors(solved_tcp, profile.initial_tcp)
            configured_tcp = physical_kinematics.forward(
                profile.configured_initial_q_rad, profile.tool_from_flange
            )
            configured_error = pose_errors(configured_tcp, profile.initial_tcp)
            q = result.q.copy()
            q.setflags(write=False)
            self.initial_pose_diagnostics[arm] = InitialPoseDiagnostic(
                arm=arm,
                solved_q_rad=q,
                solved_position_error_m=solved_error[0],
                solved_orientation_error_rad=solved_error[1],
                configured_joint_position_error_m=configured_error[0],
                configured_joint_orientation_error_rad=configured_error[1],
            )
        self.scene.set_dual_configuration(solved["left"], solved["right"])
        self._initial_joint_positions = {
            arm: solved[arm].copy() for arm in ("left", "right")
        }
        for arm in ("left", "right"):
            self.set_target_tcp_pose(arm, self.profiles[arm].initial_tcp)

    def connect(self) -> None:
        if self._closed:
            raise RuntimeError("MujocoRobot is closed")

    def read_joint_positions(self) -> DualArmJointState:
        self._require_open()
        return DualArmJointState(
            {
                arm: self.scene.get_arm_configuration(arm)
                for arm in ("left", "right")
            }
        )

    def command_joint_positions(
        self, targets_rad: Mapping[ArmName, np.ndarray]
    ) -> None:
        self._require_open()
        if not targets_rad or not set(targets_rad).issubset({"left", "right"}):
            raise ValueError("joint targets must contain at least one valid arm")
        current = self.read_joint_positions().positions_rad
        left = np.asarray(targets_rad.get("left", current["left"]), dtype=float)
        right = np.asarray(targets_rad.get("right", current["right"]), dtype=float)
        self.scene.set_dual_configuration(left, right)
        self._stopped_arms.difference_update(targets_rad)

    def set_target_tcp_pose(self, arm: ArmName, target_tcp: pin.SE3) -> None:
        self._require_open()
        if arm not in ("left", "right"):
            raise ValueError("arm must be 'left' or 'right'")
        validate_se3(target_tcp, "target_tcp")
        mount = (
            self.assembly.mounts.left_base
            if arm == "left"
            else self.assembly.mounts.right_base
        )
        self.scene.set_arm_target_marker(arm, mount * target_tcp)

    def get_tcp_pose(self, arm: ArmName) -> pin.SE3:
        self._require_open()
        q = self.scene.get_arm_configuration(arm)
        return self._kinematics[arm].forward(
            q, self.profiles[arm].tool_from_flange
        )

    def reset_to_initial(self) -> None:
        """Restore the solved startup state without recreating the MuJoCo model."""

        self._require_open()
        if self._initial_joint_positions is None:
            raise RuntimeError("MuJoCo initial joint positions are unavailable")
        viewer_lock = self._viewer.lock() if self._viewer is not None else nullcontext()
        with viewer_lock:
            self.scene.set_dual_configuration(
                self._initial_joint_positions["left"],
                self._initial_joint_positions["right"],
            )
            self.scene.data.qvel[:] = 0.0
            for arm in ("left", "right"):
                self.set_target_tcp_pose(arm, self.profiles[arm].initial_tcp)
        self._stopped_arms.update(("left", "right"))

    def stop(self, arms: Collection[ArmName] = ("left", "right")) -> None:
        self._require_open()
        selected = set(arms)
        if not selected.issubset({"left", "right"}):
            raise ValueError("stop arms must contain only left/right")
        self.scene.data.qvel[:] = 0.0
        self._stopped_arms.update(selected)

    def open_viewer(self, key_callback: Optional[Callable[[int], None]] = None):
        self._require_open()
        if self._viewer is None:
            import mujoco.viewer

            self._viewer = mujoco.viewer.launch_passive(
                self.scene.model,
                self.scene.data,
                key_callback=key_callback,
            )
            self._viewer.set_texts(
                (
                    None,
                    mujoco.mjtGridPos.mjGRID_BOTTOMLEFT,
                    "Reset dual arms",
                    "R / Home",
                )
            )
        return self._viewer

    def sync_viewer(self) -> bool:
        if self._viewer is None:
            return True
        if not self._viewer.is_running():
            return False
        self._viewer.sync()
        return True

    def render(self, width: int = 1280, height: int = 720) -> np.ndarray:
        self._require_open()
        return self.scene.render(width, height)

    def close(self) -> None:
        if self._closed:
            return
        if self._viewer is not None:
            self._viewer.close()
            self._viewer = None
            sleep(0.2)
        self.scene.data.qvel[:] = 0.0
        self._closed = True

    def _require_open(self) -> None:
        if self._closed:
            raise RuntimeError("MujocoRobot is closed")
