"""MyArmSTS: pymycobot-like API backed by direct STS control."""

from __future__ import annotations

from dataclasses import dataclass

from .constants import DEFAULT_BAUDRATE, DEFAULT_TIMEOUT, MemAddr, StatusBit
from .conversion import (
    angle_to_encoder,
    current_raw_to_ma,
    encoder_to_angle,
    voltage_raw_to_v,
)
from .exceptions import InvalidParameterError, NotImplementedFeatureError
from .joint_config import DEFAULT_ARM_CONFIG, ArmConfig
from .protocol import bytes_to_int16
from .servo_controller import ServoController
from .sts import STS

# Wide sync read layout: 15 bytes
# from PRESENT_POSITION (0x38) through PRESENT_CURRENT (0x46)
_SNAPSHOT_ADDR = MemAddr.PRESENT_POSITION  # 0x38
_SNAPSHOT_LEN = 15
_OFF_POSITION = 0
_OFF_SPEED = 2
_OFF_LOAD = 4
_OFF_VOLTAGE = 6
_OFF_TEMPERATURE = 7
_OFF_STATUS = 9
_OFF_MOVING = 10
_OFF_CURRENT = 13


@dataclass
class ServoSnapshot:
    """All status fields for one servo, captured in a single sync read."""

    servo_id: int
    encoder: int | None = None
    speed: int | None = None
    load: int | None = None
    voltage: float | None = None
    temperature: int | None = None
    status: int | None = None
    moving: bool | None = None
    current: float | None = None
    error_code: int | None = None
    responded: bool = False


class MyArmSTS:
    """High-level arm API compatible with common pymycobot usage."""

    def __init__(
        self,
        port: str,
        baudrate: int | str = DEFAULT_BAUDRATE,
        timeout: float = DEFAULT_TIMEOUT,
        config: ArmConfig | None = None,
        debug: bool = False,
    ) -> None:
        if isinstance(baudrate, bool) or not isinstance(baudrate, (int, str)):
            raise InvalidParameterError(
                f"baudrate must be int or str, got {type(baudrate).__name__}"
            )
        try:
            baudrate_int = int(baudrate)
        except ValueError:
            raise InvalidParameterError(
                f"Invalid baudrate string: {baudrate!r}"
            ) from None

        self._config = config or DEFAULT_ARM_CONFIG
        self._controller = ServoController(
            port=port,
            baudrate=baudrate_int,
            timeout=timeout,
            debug=debug,
        )
        self._sts = STS(self._controller)
        self.connect()

    # ---------------- connection lifecycle ----------------

    def connect(self) -> bool:
        """Open serial connection."""
        return self._controller.connect()

    def disconnect(self) -> None:
        """Disable all servos then close serial connection."""
        self._controller.disconnect(
            disable_ids=self._config.all_servo_ids(include_gripper=True)
        )

    def is_connected(self) -> bool:
        """Return connection state."""
        return self._controller.is_connected()

    def get_comm_stats(self) -> dict:
        """Get communication stats for runtime monitoring."""
        return self._controller.get_comm_stats()

    def print_comm_stats(self) -> None:
        """Print communication stats snapshot."""
        stats = self.get_comm_stats()
        print(
            "[comm-stats] "
            f"no_response={stats['no_response_count']}, "
            f"duration={stats['connected_duration_sec']:.3f}s, "
            f"miss_rate={stats['no_response_per_minute']:.3f}/min"
        )

    def reset_comm_stats(self) -> None:
        """Reset communication stats counters."""
        self._controller.reset_comm_stats()

    def __enter__(self) -> "MyArmSTS":
        self.connect()
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> None:
        self.disconnect()

    # ---------------- device info ----------------

    def get_serial_number(self) -> str | None:
        """Return the USB serial number of the servo adapter board.

        Uses pyserial's ``list_ports`` to match the configured port and
        read the USB iSerial descriptor.  Works on Linux, macOS, and Windows.
        Returns ``None`` if the serial number cannot be determined.
        """
        from .port import get_port_serial_number

        return get_port_serial_number(self._controller.port)

    # ---------------- encoder-level APIs ----------------

    def get_servo_encoder(self, servo_id: int) -> int:
        """Get one servo encoder value."""
        return self._sts.read_position(servo_id)

    def get_servos_encoder(self) -> list[int]:
        """Get all servo encoders as an 8-length list (servo IDs 1..8 order)."""
        motor_ids = self._config.all_servo_ids(include_gripper=True)
        values = self._sts.read_positions(motor_ids)
        return [values.get(motor_id, 0) for motor_id in motor_ids]

    def set_servo_encoder(
        self, servo_id: int, encoder: int, speed: int = 60, acceleration: int = 150
    ) -> bool:
        """Set one servo position by encoder target.

        Args:
            speed: STS speed register value (0-32767). 0 = unlimited.
                   Typical range for STS3215 12V: 1-61, STS3250: 1-103.
        """
        return self._sts.write_position(
            servo_id,
            encoder=max(0, min(4095, encoder)),
            speed=speed,
            acceleration=acceleration,
        )

    def set_servos_encoder(
        self, encoders: list[int], speed: int = 60, acceleration: int = 150
    ) -> bool:
        """Set all servo positions by encoder targets (must be 8 values).

        Args:
            speed: STS speed register value (0-32767). 0 = unlimited.
        """
        all_ids = self._config.all_servo_ids(include_gripper=True)
        expected_len = len(all_ids)
        if len(encoders) != expected_len:
            raise InvalidParameterError(
                f"Expected {expected_len} encoder values, got {len(encoders)}"
            )

        motor_targets = []
        for servo_id, encoder in zip(all_ids, encoders):
            encoder = max(0, min(4095, int(encoder)))
            motor_targets.append((servo_id, encoder))

        return self._sts.write_positions_sync(
            motor_targets,
            speed=speed,
            acceleration=acceleration,
        )

    # ---------------- angle-level APIs ----------------

    def get_joint_angle(self, joint_id: int) -> float:
        """Get one joint angle in degrees."""
        joint = self._config.get_joint(joint_id)
        if joint is None:
            raise InvalidParameterError(f"Invalid joint_id: {joint_id}")

        encoder = self.get_servo_encoder(joint.primary_motor_id())
        return encoder_to_angle(
            encoder,
            offset=joint.encoder_offset,
            direction=joint.direction,
        )

    def get_joints_angle(self) -> list[float]:
        """Get all joint angles in degrees (single sync read)."""
        encoders = self.get_servos_encoder()
        all_ids = self._config.all_servo_ids(include_gripper=True)
        encoder_by_id = dict(zip(all_ids, encoders))
        angles: list[float] = []
        for joint_id in range(1, self._config.dof + 1):
            joint = self._config.get_joint(joint_id)
            enc = encoder_by_id.get(joint.primary_motor_id(), 0)
            angles.append(
                encoder_to_angle(
                    enc,
                    offset=joint.encoder_offset,
                    direction=joint.direction,
                )
            )
        return angles

    def set_joint_angle(self, joint_id: int, angle: float, speed: int = 30) -> bool:
        """Set one joint target angle.

        Args:
            speed: STS speed register value (0-32767). 0 = unlimited.
        """
        joint = self._config.get_joint(joint_id)
        if joint is None:
            raise InvalidParameterError(f"Invalid joint_id: {joint_id}")

        bounded = max(joint.min_angle, min(joint.max_angle, angle))
        encoder = angle_to_encoder(
            bounded,
            offset=joint.encoder_offset,
            direction=joint.direction,
        )

        motor_targets = self._config.joint_to_motor_pairs(joint_id, encoder)
        return self._sts.write_positions_sync(
            motor_targets,
            speed=speed,
            acceleration=0,
        )

    def set_joints_angle(self, angles: list[float], speed: int = 30) -> bool:
        """Set all joint target angles."""
        if len(angles) != self._config.dof:
            raise InvalidParameterError(
                f"Expected {self._config.dof} angle values, got {len(angles)}"
            )

        # Keep full 8-servo command shape and only update the 6-DOF joints.
        all_servo_ids = self._config.all_servo_ids(include_gripper=True)
        current_by_id = dict(zip(all_servo_ids, self.get_servos_encoder()))

        for joint_id, angle in enumerate(angles, start=1):
            joint = self._config.get_joint(joint_id)
            if joint is None:
                raise InvalidParameterError(f"Invalid joint_id: {joint_id}")
            bounded = max(joint.min_angle, min(joint.max_angle, angle))
            joint_encoder = angle_to_encoder(
                bounded,
                offset=joint.encoder_offset,
                direction=joint.direction,
            )
            for servo_id, motor_encoder in self._config.joint_to_motor_pairs(
                joint_id, joint_encoder
            ):
                current_by_id[servo_id] = motor_encoder

        full_encoders = [current_by_id.get(servo_id, 0) for servo_id in all_servo_ids]
        return self.set_servos_encoder(full_encoders, speed=speed)

    # ---------------- state/status APIs ----------------

    def get_servos_snapshot(self) -> list[ServoSnapshot]:
        """Get all servo status in a single wide sync read (15 bytes from 0x38).

        Returns a list of ServoSnapshot, one per servo in config order.
        Fields invalidated by specific error bits are set to None.
        """
        servo_ids = self._config.all_servo_ids(include_gripper=True)
        parsed, errors = self._sts.read_registers_sync(
            servo_ids, _SNAPSHOT_ADDR, _SNAPSHOT_LEN
        )
        snapshots: list[ServoSnapshot] = []
        for sid in servo_ids:
            payload = parsed.get(sid)
            err = errors.get(sid)
            if payload is None or len(payload) < _SNAPSHOT_LEN:
                snapshots.append(
                    ServoSnapshot(
                        servo_id=sid,
                        error_code=err,
                        responded=payload is not None,
                    )
                )
                continue

            has_encoder_err = err is not None and err & StatusBit.ENCODER_ERROR
            has_current_err = err is not None and err & StatusBit.CURRENT_ERROR

            snapshots.append(
                ServoSnapshot(
                    servo_id=sid,
                    encoder=None
                    if has_encoder_err
                    else bytes_to_int16(
                        payload[_OFF_POSITION], payload[_OFF_POSITION + 1]
                    ),
                    speed=bytes_to_int16(payload[_OFF_SPEED], payload[_OFF_SPEED + 1]),
                    load=bytes_to_int16(payload[_OFF_LOAD], payload[_OFF_LOAD + 1]),
                    voltage=voltage_raw_to_v(payload[_OFF_VOLTAGE]),
                    temperature=payload[_OFF_TEMPERATURE],
                    status=payload[_OFF_STATUS],
                    moving=bool(payload[_OFF_MOVING] & 0x01),
                    current=None
                    if has_current_err
                    else current_raw_to_ma(
                        payload[_OFF_CURRENT] | (payload[_OFF_CURRENT + 1] << 8)
                    ),
                    error_code=err,
                    responded=True,
                )
            )
        return snapshots

    def is_robot_moving(self) -> int:
        """Return 1 if any motor is moving, otherwise 0 (single sync read).

        If a servo fails to respond or returns an empty/corrupt payload,
        it is conservatively treated as still moving to avoid prematurely
        ending motion waits.
        """
        servo_ids = self._config.all_servo_ids(include_gripper=False)
        parsed, _ = self._sts.read_registers_sync(
            servo_ids, MemAddr.MOVING_FLAG, 1
        )
        for servo_id in servo_ids:
            payload = parsed.get(servo_id)
            if not payload:
                return 1  # dropout or empty → assume still moving
            if payload[0] & 0x01:
                return 1
        return 0

    def stop_robot(self) -> bool:
        """Soft stop by commanding current position as target."""
        current = self.get_servos_encoder()
        return self.set_servos_encoder(current, speed=0)

    def get_servos_temp(self) -> list[int]:
        """Get servo temperatures in Celsius (single sync read)."""
        servo_ids = self._config.all_servo_ids(include_gripper=True)
        parsed, _ = self._sts.read_registers_sync(
            servo_ids, MemAddr.PRESENT_TEMPERATURE, 1
        )
        return [parsed.get(sid, [0])[0] if parsed.get(sid) else 0 for sid in servo_ids]

    def get_servos_voltage(self) -> list[float]:
        """Get servo voltages in volts (single sync read)."""
        servo_ids = self._config.all_servo_ids(include_gripper=True)
        parsed, _ = self._sts.read_registers_sync(
            servo_ids, MemAddr.PRESENT_VOLTAGE, 1
        )
        return [
            voltage_raw_to_v(parsed[sid][0]) if parsed.get(sid) else 0.0
            for sid in servo_ids
        ]

    def get_servos_current(self) -> list[float]:
        """Get servo currents in mA (single sync read)."""
        servo_ids = self._config.all_servo_ids(include_gripper=True)
        parsed, _ = self._sts.read_registers_sync(
            servo_ids, MemAddr.PRESENT_CURRENT, 2
        )
        result: list[float] = []
        for sid in servo_ids:
            payload = parsed.get(sid)
            if payload and len(payload) >= 2:
                result.append(current_raw_to_ma(payload[0] | (payload[1] << 8)))
            else:
                result.append(0.0)
        return result

    def get_robot_error_status(self, prefer_cached: bool = False) -> list[int]:
        """Get servo status bytes.

        Args:
            prefer_cached: Use cached response-frame error code first, and
                fall back to direct STATUS register read when missing.
        """
        servo_ids = self._config.all_servo_ids(include_gripper=True)
        if prefer_cached:
            statuses: list[int] = []
            need_read: list[int] = []
            for servo_id in servo_ids:
                cached = self._sts.get_cached_error_code(servo_id)
                if cached is not None:
                    statuses.append(cached)
                else:
                    statuses.append(-1)  # placeholder
                    need_read.append(servo_id)
            if need_read:
                parsed, _ = self._sts.read_registers_sync(
                    need_read, MemAddr.STATUS, 1
                )
                for i, servo_id in enumerate(servo_ids):
                    if statuses[i] == -1:
                        payload = parsed.get(servo_id)
                        statuses[i] = payload[0] if payload else 0xFF
            return statuses

        parsed, _ = self._sts.read_registers_sync(servo_ids, MemAddr.STATUS, 1)
        return [parsed[sid][0] if parsed.get(sid) else 0xFF for sid in servo_ids]

    def is_all_servos_enabled(self) -> list[int]:
        """Get torque enable state (1/0) for all servos (single sync read)."""
        servo_ids = self._config.all_servo_ids(include_gripper=True)
        parsed, _ = self._sts.read_registers_sync(
            servo_ids, MemAddr.TORQUE_ENABLE, 1
        )
        return [
            1 if parsed.get(sid) and parsed[sid][0] == 1 else 0 for sid in servo_ids
        ]

    # ---------------- calibration/release APIs ----------------

    def set_servo_enabled(self, servo_id: int, status: int | bool) -> bool:
        """Enable/disable one servo torque."""
        return self._sts.set_torque_enabled(servo_id, bool(status))

    def set_servo_calibrate(self, servo_id: int) -> bool:
        """Calibrate one servo current pose as encoder center 2048."""
        return self._sts.calibrate_current_as_center(servo_id)

    def set_servo_offset_calibrate(
        self, servo_id: int, target_position: int
    ) -> bool:
        """Calibrate one servo: set current physical position as *target_position*.

        Sends POSITION_CALIBRATE (0x0B) instruction. The servo auto-computes
        and stores the offset in EEPROM. Requires STS firmware >= 3.10.

        Args:
            servo_id: Servo ID (1-8).
            target_position: Target encoder value for the current position (0-4095).
        Returns:
            True if servo acknowledged with no error.
        """
        return self._sts.calibrate_current_as_position(servo_id, target_position)

    def get_servo_offset(self, servo_id: int) -> int:
        """Read raw offset register (0x1F) for diagnostics.

        Returns raw 16-bit value (STS sign-magnitude encoding, 0-8191).
        """
        return self._sts.read_offset_position(servo_id)

    def release_joint(self, joint_id: int, include_dual_motor: bool = True) -> bool:
        """Release one joint so it can be moved by hand."""
        joint = self._config.get_joint(joint_id)
        if joint is None:
            raise InvalidParameterError(f"Invalid joint_id: {joint_id}")

        motor_ids = (
            joint.motor_ids if include_dual_motor else [joint.primary_motor_id()]
        )
        ok = True
        for servo_id in motor_ids:
            ok = self.set_servo_enabled(servo_id, 0) and ok
        return ok

    def focus_joint(self, joint_id: int, include_dual_motor: bool = True) -> bool:
        """Enable one joint torque to hold pose."""
        joint = self._config.get_joint(joint_id)
        if joint is None:
            raise InvalidParameterError(f"Invalid joint_id: {joint_id}")

        motor_ids = (
            joint.motor_ids if include_dual_motor else [joint.primary_motor_id()]
        )
        ok = True
        for servo_id in motor_ids:
            ok = self.set_servo_enabled(servo_id, 1) and ok
        return ok

    def release_all_joints(self, include_dual_motor: bool = True) -> bool:
        """Release all joints and gripper."""
        ok = True
        for joint_id in range(1, self._config.dof + 1):
            ok = (
                self.release_joint(joint_id, include_dual_motor=include_dual_motor)
                and ok
            )
        if self._config.gripper is not None:
            ok = self.set_servo_enabled(self._config.gripper.motor_id, 0) and ok
        return ok

    def focus_all_joints(self, include_dual_motor: bool = True) -> bool:
        """Enable all joints and gripper."""
        ok = True
        for joint_id in range(1, self._config.dof + 1):
            ok = (
                self.focus_joint(joint_id, include_dual_motor=include_dual_motor) and ok
            )
        if self._config.gripper is not None:
            ok = self.set_servo_enabled(self._config.gripper.motor_id, 1) and ok
        return ok

    # ---------------- gripper APIs ----------------

    def set_gripper_value(self, value: int, speed: int = 30) -> bool:
        """Set gripper opening percentage [0..100]."""
        if self._config.gripper is None:
            raise NotImplementedFeatureError("gripper")
        value = int(max(0, min(100, value)))
        gripper = self._config.gripper
        encoder = gripper.open_encoder + int(
            (gripper.close_encoder - gripper.open_encoder) * value / 100
        )
        return self.set_servo_encoder(gripper.motor_id, encoder, speed=speed)

    def get_gripper_value(self) -> int:
        """Get gripper opening percentage [0..100]."""
        if self._config.gripper is None:
            raise NotImplementedFeatureError("gripper")
        gripper = self._config.gripper
        encoder = self.get_servo_encoder(gripper.motor_id)
        span = max(1, gripper.close_encoder - gripper.open_encoder)
        value = int((encoder - gripper.open_encoder) * 100 / span)
        return int(max(0, min(100, value)))

    # ---------------- explicit unsupported APIs ----------------

    def get_joints_coord(self) -> None:
        raise NotImplementedFeatureError("get_joints_coord")

    def set_joints_coord(self, coords: list[float], speed: int = 30) -> None:
        raise NotImplementedFeatureError("set_joints_coord")

    def power_on(self) -> None:
        raise NotImplementedFeatureError("power_on")

    def power_off(self) -> None:
        raise NotImplementedFeatureError("power_off")

    def get_power_status(self) -> None:
        raise NotImplementedFeatureError("get_power_status")

    def set_master_out_io_state(self, pin: int, state: int) -> None:
        raise NotImplementedFeatureError("set_master_out_io_state")

    def get_master_in_io_state(self, pin: int) -> None:
        raise NotImplementedFeatureError("get_master_in_io_state")

    def set_tool_out_io_state(self, pin: int, state: int) -> None:
        raise NotImplementedFeatureError("set_tool_out_io_state")

    def get_tool_in_io_state(self, pin: int) -> None:
        raise NotImplementedFeatureError("get_tool_in_io_state")

    def set_tool_led_color(self, r: int, g: int, b: int) -> None:
        raise NotImplementedFeatureError("set_tool_led_color")

    def is_tool_btn_clicked(self, mode: int = 1) -> None:
        raise NotImplementedFeatureError("is_tool_btn_clicked")
