from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Dict, Mapping, Optional

import numpy as np
import pinocchio as pin

from .kinematics import RM75Kinematics, validate_se3
from .robot_backend import RobotBackend
from .solver import RM75IkSolver
from .teleop_config import ArmName, ArmTeleopProfile
from .types import IkOptions, teleop_joint_limits


class SafetyState(str, Enum):
    IDLE = "idle"
    ACTIVE = "active"
    RESETTING = "resetting"
    FAULT = "fault"


@dataclass(frozen=True)
class ControllerSample:
    hand: ArmName
    grip: bool
    trigger: float
    position_m: np.ndarray
    quaternion_xyzw: np.ndarray

    def __post_init__(self) -> None:
        if self.hand not in ("left", "right"):
            raise ValueError("controller hand must be 'left' or 'right'")
        position = np.asarray(self.position_m, dtype=float).copy()
        quaternion = np.asarray(self.quaternion_xyzw, dtype=float).copy()
        if position.shape != (3,) or not np.all(np.isfinite(position)):
            raise ValueError("controller position must be finite with shape (3,)")
        if quaternion.shape != (4,) or not np.all(np.isfinite(quaternion)):
            raise ValueError("controller quaternion must be finite with shape (4,)")
        norm = float(np.linalg.norm(quaternion))
        if norm <= 1e-9:
            raise ValueError("controller quaternion has zero norm")
        quaternion /= norm
        if not np.isfinite(self.trigger):
            raise ValueError("controller trigger must be finite")
        position.setflags(write=False)
        quaternion.setflags(write=False)
        object.__setattr__(self, "position_m", position)
        object.__setattr__(self, "quaternion_xyzw", quaternion)

    @classmethod
    def from_message(
        cls, message, expected_arm: Optional[ArmName] = None
    ) -> "ControllerSample":
        message_hand = str(getattr(message, "hand", "")).strip().lower()
        hand = expected_arm or message_hand
        if expected_arm is not None and message_hand and message_hand != expected_arm:
            raise ValueError(
                f"controller message hand {message_hand!r} does not match {expected_arm!r}"
            )
        pose = message.pose
        return cls(
            hand=hand,
            grip=bool(message.grip),
            trigger=float(message.trigger),
            position_m=np.array(
                [pose.position.x, pose.position.y, pose.position.z], dtype=float
            ),
            quaternion_xyzw=np.array(
                [
                    pose.orientation.x,
                    pose.orientation.y,
                    pose.orientation.z,
                    pose.orientation.w,
                ],
                dtype=float,
            ),
        )


@dataclass(frozen=True)
class MappedTarget:
    target_tcp: pin.SE3
    clamped: bool
    just_engaged: bool


@dataclass(frozen=True)
class ControlCycleResult:
    state: SafetyState
    commanded_arms: tuple[ArmName, ...] = ()
    targets_tcp: Optional[Mapping[ArmName, pin.SE3]] = None
    reason: str = ""


class RelativePoseMapper:
    def __init__(self, profile: ArmTeleopProfile) -> None:
        self.profile = profile
        self.active = False
        self._controller_start_position: Optional[np.ndarray] = None
        self._controller_start_rotation: Optional[np.ndarray] = None
        self._robot_start_tcp: Optional[pin.SE3] = None
        self._filtered_position: Optional[np.ndarray] = None
        self._filtered_rotation: Optional[np.ndarray] = None
        self._last_position: Optional[np.ndarray] = None
        self._last_rotation: Optional[np.ndarray] = None

    def reset(self) -> None:
        self.active = False
        self._controller_start_position = None
        self._controller_start_rotation = None
        self._robot_start_tcp = None
        self._filtered_position = None
        self._filtered_rotation = None
        self._last_position = None
        self._last_rotation = None

    @staticmethod
    def _rotation(sample: ControllerSample) -> np.ndarray:
        x, y, z, w = sample.quaternion_xyzw
        return pin.Quaternion(w, x, y, z).matrix()

    def map(
        self,
        sample: ControllerSample,
        current_tcp: pin.SE3,
        dt: float,
    ) -> MappedTarget:
        if sample.hand != self.profile.arm:
            raise ValueError("controller sample was routed to the wrong arm")
        validate_se3(current_tcp, "current_tcp")
        if not np.isfinite(dt) or dt <= 0.0:
            raise ValueError("control dt must be finite and positive")
        rotation = self._rotation(sample)
        if not self.active:
            self.active = True
            self._controller_start_position = sample.position_m.copy()
            self._controller_start_rotation = rotation.copy()
            self._robot_start_tcp = current_tcp.copy()
            self._filtered_position = current_tcp.translation.copy()
            self._filtered_rotation = current_tcp.rotation.copy()
            self._last_position = current_tcp.translation.copy()
            self._last_rotation = current_tcp.rotation.copy()
            return MappedTarget(current_tcp.copy(), False, True)

        assert self._controller_start_position is not None
        assert self._controller_start_rotation is not None
        assert self._robot_start_tcp is not None
        delta = sample.position_m - self._controller_start_position
        mapped_delta = self.profile.xr_to_robot @ delta
        raw_position = (
            self._robot_start_tcp.translation + self.profile.scale * mapped_delta
        )
        raw_position = np.where(
            np.asarray(self.profile.enable_position_axes, dtype=bool),
            raw_position,
            self._robot_start_tcp.translation,
        )
        position, clamped = self._clamp_workspace(raw_position)
        position = self._filter_position(position)
        position, limited = self._limit_position_step(position, dt)
        position, final_clamped = self._clamp_workspace(position)

        target_rotation = self._target_rotation(rotation)
        target_rotation = self._filter_rotation(target_rotation)
        target_rotation, rotation_limited = self._limit_rotation_step(
            target_rotation, dt
        )
        self._last_position = position.copy()
        self._last_rotation = target_rotation.copy()
        return MappedTarget(
            pin.SE3(target_rotation, position),
            clamped or limited or final_clamped or rotation_limited,
            False,
        )

    def _clamp_workspace(self, target: np.ndarray) -> tuple[np.ndarray, bool]:
        result = np.clip(
            np.asarray(target, dtype=float),
            self.profile.workspace_min,
            self.profile.workspace_max,
        )
        min_radius, max_radius = self.profile.cylinder_radius_limit
        if result[2] < self.profile.low_z_threshold:
            min_radius = max(min_radius, self.profile.low_z_min_radius)
        radius = float(np.hypot(result[0], result[1]))
        if radius > max_radius:
            result[:2] *= max_radius / radius
        elif radius < min_radius:
            if radius > 1e-9:
                result[:2] *= min_radius / radius
            else:
                result[:2] = [min_radius, 0.0]
        changed = not np.allclose(result, target, atol=1e-12, rtol=0.0)
        return result, changed

    def _filter_position(self, target: np.ndarray) -> np.ndarray:
        assert self._filtered_position is not None
        assert self._last_position is not None
        if np.linalg.norm(target - self._last_position) < self.profile.deadband_m:
            target = self._last_position
        distance = float(np.linalg.norm(target - self._filtered_position))
        ratio = min(
            1.0, distance / self.profile.target_filter_fast_threshold_m
        )
        alpha = self.profile.target_filter_alpha + ratio * (
            self.profile.target_filter_alpha_fast
            - self.profile.target_filter_alpha
        )
        self._filtered_position = (
            alpha * target + (1.0 - alpha) * self._filtered_position
        )
        return self._filtered_position.copy()

    def _limit_position_step(
        self, target: np.ndarray, dt: float
    ) -> tuple[np.ndarray, bool]:
        assert self._last_position is not None
        delta = target - self._last_position
        distance = float(np.linalg.norm(delta))
        maximum = self.profile.max_linear_speed_m_s * dt
        if distance <= maximum or distance <= 1e-12:
            return target, False
        return self._last_position + delta * (maximum / distance), True

    def _target_rotation(self, controller_rotation: np.ndarray) -> np.ndarray:
        assert self._controller_start_rotation is not None
        assert self._robot_start_tcp is not None
        if not self.profile.enable_orientation_control:
            return self._robot_start_tcp.rotation.copy()
        xr_delta = controller_rotation @ self._controller_start_rotation.T
        matrix = self.profile.xr_to_robot
        robot_delta = matrix @ xr_delta @ matrix.T
        target = robot_delta @ self._robot_start_tcp.rotation
        axes = np.asarray(self.profile.enable_orientation_axes, dtype=bool)
        if not np.all(axes):
            target_rpy = pin.rpy.matrixToRpy(target)
            start_rpy = pin.rpy.matrixToRpy(self._robot_start_tcp.rotation)
            target = pin.rpy.rpyToMatrix(*np.where(axes, target_rpy, start_rpy))
        return target

    def _filter_rotation(self, target: np.ndarray) -> np.ndarray:
        assert self._filtered_rotation is not None
        assert self._last_rotation is not None
        if (
            np.linalg.norm(pin.log3(self._last_rotation.T @ target))
            < self.profile.orientation_deadband_rad
        ):
            target = self._last_rotation
        delta = pin.log3(self._filtered_rotation.T @ target)
        self._filtered_rotation = self._filtered_rotation @ pin.exp3(
            self.profile.orientation_filter_alpha * delta
        )
        return self._filtered_rotation.copy()

    def _limit_rotation_step(
        self, target: np.ndarray, dt: float
    ) -> tuple[np.ndarray, bool]:
        assert self._last_rotation is not None
        delta = pin.log3(self._last_rotation.T @ target)
        angle = float(np.linalg.norm(delta))
        maximum = self.profile.max_orientation_speed_rad_s * dt
        if angle <= maximum or angle <= 1e-12:
            return target, False
        return self._last_rotation @ pin.exp3(delta * (maximum / angle)), True


class DualArmQpTeleopController:
    def __init__(
        self,
        robot: RobotBackend,
        profiles: Mapping[ArmName, ArmTeleopProfile],
        control_rate_hz: float = 90.0,
        ik_options: Optional[IkOptions] = None,
    ) -> None:
        if set(profiles) != {"left", "right"}:
            raise ValueError("profiles must contain left and right arms")
        if not np.isfinite(control_rate_hz) or control_rate_hz <= 0.0:
            raise ValueError("control_rate_hz must be finite and positive")
        self.robot = robot
        self.profiles = dict(profiles)
        self.dt = 1.0 / control_rate_hz
        self.ik_options = ik_options or IkOptions(
            max_iterations=120,
            time_limit_sec=0.008,
        )
        self.kinematics = {
            arm: RM75Kinematics(limits=teleop_joint_limits())
            for arm in ("left", "right")
        }
        self.solvers = {
            arm: RM75IkSolver(self.kinematics[arm]) for arm in ("left", "right")
        }
        self.mappers = {
            arm: RelativePoseMapper(self.profiles[arm])
            for arm in ("left", "right")
        }
        self._latest: Dict[ArmName, tuple[ControllerSample, float]] = {}
        self._fault_reason = ""
        self._reset_waiting_for_release = False
        self._closed = False
        self.robot.connect()

    @property
    def safety_state(self) -> SafetyState:
        if self._fault_reason:
            return SafetyState.FAULT
        if self._reset_waiting_for_release:
            return SafetyState.RESETTING
        if any(mapper.active for mapper in self.mappers.values()):
            return SafetyState.ACTIVE
        return SafetyState.IDLE

    def update_controller(
        self,
        message,
        timestamp_sec: float,
        expected_arm: Optional[ArmName] = None,
    ) -> None:
        self.update_sample(
            ControllerSample.from_message(message, expected_arm), timestamp_sec
        )

    def update_sample(self, sample: ControllerSample, timestamp_sec: float) -> None:
        if not np.isfinite(timestamp_sec):
            raise ValueError("controller timestamp must be finite")
        self._latest[sample.hand] = (sample, float(timestamp_sec))

    def reject_input(self, reason: str) -> None:
        self._trip_fault(f"invalid controller input: {reason}")

    def reset_to_initial(self) -> ControlCycleResult:
        """Reset the backend and require fresh dual-grip release before rearming."""

        if self._closed:
            raise RuntimeError("controller is closed")
        for mapper in self.mappers.values():
            mapper.reset()
        try:
            self.robot.stop(("left", "right"))
            self.robot.reset_to_initial()
        except Exception as exc:
            self._trip_fault(f"reset/backend failure: {exc}")
            raise RuntimeError(f"failed to reset robot backend: {exc}") from exc
        self._latest.clear()
        self._fault_reason = ""
        self._reset_waiting_for_release = True
        return ControlCycleResult(
            SafetyState.RESETTING,
            reason="reset complete; waiting for fresh dual-grip release",
        )

    def step(self, timestamp_sec: float) -> ControlCycleResult:
        if self._closed:
            raise RuntimeError("controller is closed")
        if not np.isfinite(timestamp_sec):
            return self._trip_fault("control timestamp is non-finite")
        if self._fault_reason:
            if self._can_clear_fault(timestamp_sec):
                self._fault_reason = ""
                return ControlCycleResult(SafetyState.IDLE, reason="fault cleared")
            return ControlCycleResult(SafetyState.FAULT, reason=self._fault_reason)
        if self._reset_waiting_for_release:
            if self._can_finish_reset(timestamp_sec):
                self._reset_waiting_for_release = False
                self._latest.clear()
                return ControlCycleResult(SafetyState.IDLE, reason="reset rearmed")
            return ControlCycleResult(
                SafetyState.RESETTING,
                reason="waiting for fresh dual-grip release",
            )

        try:
            state = self.robot.read_joint_positions()
            commands: Dict[ArmName, np.ndarray] = {}
            targets: Dict[ArmName, pin.SE3] = {}
            for arm in ("left", "right"):
                latest = self._latest.get(arm)
                mapper = self.mappers[arm]
                if latest is None:
                    continue
                sample, sample_time = latest
                age = timestamp_sec - sample_time
                if age < -1e-6:
                    return self._trip_fault(f"{arm} controller timestamp is in the future")
                if not sample.grip:
                    if mapper.active:
                        mapper.reset()
                        self.robot.stop((arm,))
                    continue
                if age > self.profiles[arm].command_timeout_sec:
                    return self._trip_fault(f"{arm} controller input timed out")

                q_current = state.positions_rad[arm]
                current_tcp = self.kinematics[arm].forward(
                    q_current, self.profiles[arm].tool_from_flange
                )
                mapped = mapper.map(sample, current_tcp, self.dt)
                flange_target = (
                    mapped.target_tcp * self.profiles[arm].tool_from_flange.inverse()
                )
                result = self.solvers[arm].solve(
                    flange_target, q_current, self.ik_options
                )
                if not result.success or result.q is None:
                    return self._trip_fault(
                        f"{arm} IK failed: {result.status.value}: {result.message}"
                    )
                max_step = self.profiles[arm].joint_max_speed_rad_s * self.dt
                q_command = np.clip(result.q, q_current - max_step, q_current + max_step)
                limits = self.kinematics[arm].limits
                q_command = np.clip(q_command, limits.lower, limits.upper)
                commands[arm] = q_command
                targets[arm] = mapped.target_tcp

            if commands:
                self.robot.command_joint_positions(commands)
                for arm, target in targets.items():
                    self.robot.set_target_tcp_pose(arm, target)
            return ControlCycleResult(
                self.safety_state,
                tuple(commands),
                targets or None,
            )
        except Exception as exc:
            return self._trip_fault(f"control/backend failure: {exc}")

    def _can_clear_fault(self, timestamp_sec: float) -> bool:
        if set(self._latest) != {"left", "right"}:
            return False
        for arm, (sample, sample_time) in self._latest.items():
            if sample.grip:
                return False
            if timestamp_sec - sample_time > self.profiles[arm].command_timeout_sec:
                return False
        for mapper in self.mappers.values():
            mapper.reset()
        return True

    def _can_finish_reset(self, timestamp_sec: float) -> bool:
        if set(self._latest) != {"left", "right"}:
            return False
        for arm, (sample, sample_time) in self._latest.items():
            if sample.grip:
                return False
            age = timestamp_sec - sample_time
            if age < -1e-6 or age > self.profiles[arm].command_timeout_sec:
                return False
        for mapper in self.mappers.values():
            mapper.reset()
        return True

    def _trip_fault(self, reason: str) -> ControlCycleResult:
        self._fault_reason = str(reason)
        self._reset_waiting_for_release = False
        for mapper in self.mappers.values():
            mapper.reset()
        try:
            self.robot.stop(("left", "right"))
        except Exception:
            pass
        return ControlCycleResult(SafetyState.FAULT, reason=self._fault_reason)

    def stop(self) -> None:
        self._reset_waiting_for_release = False
        for mapper in self.mappers.values():
            mapper.reset()
        self.robot.stop(("left", "right"))

    def close(self) -> None:
        if self._closed:
            return
        try:
            self.stop()
        finally:
            self.robot.close()
            self._closed = True
