from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Literal, Mapping

import numpy as np
import pinocchio as pin
import yaml

from .kinematics import validate_se3


ArmName = Literal["left", "right"]


def _readonly_vector(values, shape: tuple[int, ...], name: str) -> np.ndarray:
    array = np.asarray(values, dtype=float).copy()
    if array.shape != shape or not np.all(np.isfinite(array)):
        raise ValueError(f"{name} must be finite with shape {shape}")
    array.setflags(write=False)
    return array


@dataclass(frozen=True)
class ArmTeleopProfile:
    arm: ArmName
    initial_tcp: pin.SE3
    configured_initial_q_rad: np.ndarray
    tool_from_flange: pin.SE3
    active_tool_name: str
    xr_to_robot: np.ndarray
    scale: float
    command_timeout_sec: float
    deadband_m: float
    target_filter_alpha: float
    target_filter_alpha_fast: float
    target_filter_fast_threshold_m: float
    max_linear_speed_m_s: float
    enable_position_axes: tuple[bool, bool, bool]
    enable_orientation_control: bool
    enable_orientation_axes: tuple[bool, bool, bool]
    orientation_deadband_rad: float
    orientation_filter_alpha: float
    max_orientation_speed_rad_s: float
    workspace_min: np.ndarray
    workspace_max: np.ndarray
    cylinder_radius_limit: np.ndarray
    low_z_threshold: float
    low_z_min_radius: float
    joint_max_speed_rad_s: np.ndarray

    def __post_init__(self) -> None:
        if self.arm not in ("left", "right"):
            raise ValueError("arm must be 'left' or 'right'")
        validate_se3(self.initial_tcp, "initial_tcp")
        validate_se3(self.tool_from_flange, "tool_from_flange")
        object.__setattr__(
            self,
            "configured_initial_q_rad",
            _readonly_vector(self.configured_initial_q_rad, (7,), "configured_initial_q_rad"),
        )
        matrix = _readonly_vector(self.xr_to_robot, (3, 3), "xr_to_robot")
        if not np.allclose(matrix.T @ matrix, np.eye(3), atol=1e-9):
            raise ValueError("xr_to_robot must be orthonormal")
        if not np.isclose(np.linalg.det(matrix), 1.0, atol=1e-9):
            raise ValueError("xr_to_robot must be a proper rotation")
        object.__setattr__(self, "xr_to_robot", matrix)
        workspace_min = _readonly_vector(self.workspace_min, (3,), "workspace_min")
        workspace_max = _readonly_vector(self.workspace_max, (3,), "workspace_max")
        if np.any(workspace_min >= workspace_max):
            raise ValueError("workspace_min must be below workspace_max")
        object.__setattr__(self, "workspace_min", workspace_min)
        object.__setattr__(self, "workspace_max", workspace_max)
        cylinder = _readonly_vector(
            self.cylinder_radius_limit, (2,), "cylinder_radius_limit"
        )
        if cylinder[0] < 0.0 or cylinder[1] <= cylinder[0]:
            raise ValueError("cylinder radius limits are invalid")
        object.__setattr__(self, "cylinder_radius_limit", cylinder)
        speeds = _readonly_vector(
            self.joint_max_speed_rad_s, (7,), "joint_max_speed_rad_s"
        )
        if np.any(speeds <= 0.0):
            raise ValueError("joint_max_speed_rad_s must be positive")
        object.__setattr__(self, "joint_max_speed_rad_s", speeds)
        for name in (
            "scale",
            "command_timeout_sec",
            "target_filter_fast_threshold_m",
            "max_linear_speed_m_s",
            "max_orientation_speed_rad_s",
        ):
            if not np.isfinite(getattr(self, name)) or getattr(self, name) <= 0.0:
                raise ValueError(f"{name} must be finite and positive")
        for name in (
            "deadband_m",
            "orientation_deadband_rad",
            "low_z_threshold",
            "low_z_min_radius",
        ):
            if not np.isfinite(getattr(self, name)) or getattr(self, name) < 0.0:
                raise ValueError(f"{name} must be finite and non-negative")
        for name in (
            "target_filter_alpha",
            "target_filter_alpha_fast",
            "orientation_filter_alpha",
        ):
            if not 0.0 <= getattr(self, name) <= 1.0:
                raise ValueError(f"{name} must be in [0, 1]")


def _pose_from_rpy(values, name: str) -> pin.SE3:
    pose = _readonly_vector(values, (6,), name)
    return pin.SE3(pin.rpy.rpyToMatrix(*pose[3:]), pose[:3])


def _tool_pose(values, name: str) -> pin.SE3:
    pose = _readonly_vector(values, (7,), name)
    quaternion = pin.Quaternion(pose[6], pose[3], pose[4], pose[5])
    if quaternion.norm() <= 1e-12:
        raise ValueError(f"{name} quaternion has zero norm")
    quaternion.normalize()
    return pin.SE3(quaternion.matrix(), pose[:3])


def load_dual_arm_profiles(
    teleop_config_path: Path | str,
    peripheral_config_path: Path | str,
) -> Dict[ArmName, ArmTeleopProfile]:
    teleop_path = Path(teleop_config_path)
    peripheral_path = Path(peripheral_config_path)
    with teleop_path.open("r", encoding="utf-8") as stream:
        teleop_document = yaml.safe_load(stream)
    with peripheral_path.open("r", encoding="utf-8") as stream:
        peripheral_document = yaml.safe_load(stream)
    if not isinstance(teleop_document, Mapping) or not isinstance(
        peripheral_document, Mapping
    ):
        raise ValueError("teleop and peripheral YAML documents must be mappings")

    tools = peripheral_document.get("tools_in_ee")
    arms = peripheral_document.get("arms")
    if not isinstance(tools, Mapping) or not isinstance(arms, Mapping):
        raise ValueError("peripheral configuration is missing tools_in_ee or arms")
    tool_names = list(tools)

    profiles: Dict[ArmName, ArmTeleopProfile] = {}
    for arm in ("left", "right"):
        section_name = f"{arm}_arm_teleop"
        section = teleop_document.get(section_name)
        if not isinstance(section, Mapping):
            raise ValueError(f"missing {section_name} configuration")
        params = section.get("ros__parameters")
        if not isinstance(params, Mapping):
            raise ValueError(f"{section_name} is missing ros__parameters")
        arm_tool = arms.get(arm)
        if not isinstance(arm_tool, Mapping):
            raise ValueError(f"peripheral configuration is missing arm {arm}")
        tool_index = int(arm_tool.get("scissorgripper"))
        try:
            tool_name = tool_names[tool_index]
            tool_values = tools[tool_name]["pose"]
        except (IndexError, KeyError, TypeError) as exc:
            raise ValueError(f"invalid active tool selection for {arm}") from exc

        joint_speed = float(params["joint_max_speed"])
        profiles[arm] = ArmTeleopProfile(
            arm=arm,
            initial_tcp=_pose_from_rpy(params["initial_tcp_pose"], f"{arm}.initial_tcp_pose"),
            configured_initial_q_rad=np.deg2rad(params["initial_joint_pose"]),
            tool_from_flange=_tool_pose(tool_values, f"tools.{tool_name}.pose"),
            active_tool_name=tool_name,
            xr_to_robot=np.asarray(params["xr_to_robot_matrix"], dtype=float).reshape(3, 3),
            scale=float(params["scale"]),
            command_timeout_sec=float(params["command_timeout_sec"]),
            deadband_m=float(params["deadband_m"]),
            target_filter_alpha=float(params["target_filter_alpha"]),
            target_filter_alpha_fast=float(params["target_filter_alpha_fast"]),
            target_filter_fast_threshold_m=float(
                params["target_filter_fast_threshold_m"]
            ),
            max_linear_speed_m_s=float(params["max_linear_speed"]),
            enable_position_axes=tuple(bool(value) for value in params["enable_position_axes"]),
            enable_orientation_control=bool(params["enable_orientation_control"]),
            enable_orientation_axes=tuple(bool(value) for value in params["enable_orientation_axes"]),
            orientation_deadband_rad=float(params["orientation_deadband_rad"]),
            orientation_filter_alpha=float(params["orientation_filter_alpha"]),
            max_orientation_speed_rad_s=float(params["max_orientation_speed"]),
            workspace_min=params["workspace_min"],
            workspace_max=params["workspace_max"],
            cylinder_radius_limit=params["cyl_radius_limit"],
            low_z_threshold=float(params["low_z_threshold"]),
            low_z_min_radius=float(params["low_z_min_radius"]),
            joint_max_speed_rad_s=np.full(7, np.deg2rad(joint_speed)),
        )
    return profiles
