"""Arm joint and motor mapping configuration."""

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class JointConfig:
    """Single joint mapping config."""

    joint_id: int
    motor_ids: list[int]
    is_dual_motor: bool = False
    encoder_offset: int = 2048
    direction: int = 1
    min_angle: float = -180.0
    max_angle: float = 180.0

    def primary_motor_id(self) -> int:
        """Return primary motor ID."""
        return self.motor_ids[0]


@dataclass
class GripperConfig:
    """Gripper mapping config."""

    motor_id: int = 8
    open_encoder: int = 0
    close_encoder: int = 4095


@dataclass
class ArmConfig:
    """Arm-level mapping config."""

    dof: int = 6
    joints: list[JointConfig] = field(default_factory=list)
    gripper: GripperConfig | None = None

    def get_joint(self, joint_id: int) -> JointConfig | None:
        """Get joint config by ID."""
        for joint in self.joints:
            if joint.joint_id == joint_id:
                return joint
        return None

    def all_servo_ids(self, include_gripper: bool = True) -> list[int]:
        """Get all servo IDs in arm order."""
        ids: list[int] = []
        for joint in self.joints:
            ids.extend(joint.motor_ids)
        if include_gripper and self.gripper is not None:
            ids.append(self.gripper.motor_id)
        return ids

    def joint_to_motor_pairs(
        self, joint_id: int, encoder: int
    ) -> list[tuple[int, int]]:
        """Map joint target encoder to one/two motor target pairs."""
        joint = self.get_joint(joint_id)
        if joint is None:
            raise ValueError(f"Invalid joint_id: {joint_id}")

        if joint.is_dual_motor:
            motor_a, motor_b = joint.motor_ids
            return [(motor_a, encoder), (motor_b, 4096 - encoder)]
        return [(joint.primary_motor_id(), encoder)]

    def motor_offsets(self) -> list[int]:
        """Return offsets for 1..dof joints."""
        return [self.joints[i].encoder_offset for i in range(self.dof)]

    def motor_directions(self) -> list[int]:
        """Return directions for 1..dof joints."""
        return [self.joints[i].direction for i in range(self.dof)]

    def motor_id_for_joint_angle_read(self, joint_id: int) -> int:
        """Read joint angle by primary motor ID."""
        joint = self.get_joint(joint_id)
        if joint is None:
            raise ValueError(f"Invalid joint_id: {joint_id}")
        return joint.primary_motor_id()


DEFAULT_ARM_CONFIG = ArmConfig(
    dof=6,
    joints=[
        JointConfig(1, [1], False, 2048, 1, -180.0, 180.0),
        JointConfig(2, [2, 3], True, 2048, 1, -90.0, 90.0),
        JointConfig(3, [4], False, 2048, 1, -150.0, 150.0),
        JointConfig(4, [5], False, 2048, 1, -180.0, 180.0),
        JointConfig(5, [6], False, 2048, 1, -180.0, 180.0),
        JointConfig(6, [7], False, 2048, 1, -180.0, 180.0),
    ],
    gripper=GripperConfig(motor_id=8, open_encoder=0, close_encoder=4095),
)
