from __future__ import annotations

from dataclasses import dataclass
from typing import Collection, Dict, Mapping, Protocol, runtime_checkable

import numpy as np
import pinocchio as pin

from .teleop_config import ArmName


@dataclass(frozen=True)
class DualArmJointState:
    positions_rad: Mapping[ArmName, np.ndarray]

    def __post_init__(self) -> None:
        normalized: Dict[ArmName, np.ndarray] = {}
        if set(self.positions_rad) != {"left", "right"}:
            raise ValueError("dual-arm state must contain left and right positions")
        for arm, values in self.positions_rad.items():
            q = np.asarray(values, dtype=float).copy()
            if q.shape != (7,) or not np.all(np.isfinite(q)):
                raise ValueError(f"{arm} joint state must be finite with shape (7,)")
            q.setflags(write=False)
            normalized[arm] = q
        object.__setattr__(self, "positions_rad", normalized)


@runtime_checkable
class RobotBackend(Protocol):
    def connect(self) -> None: ...

    def read_joint_positions(self) -> DualArmJointState: ...

    def command_joint_positions(
        self, targets_rad: Mapping[ArmName, np.ndarray]
    ) -> None: ...

    def set_target_tcp_pose(self, arm: ArmName, target_tcp: pin.SE3) -> None: ...

    def reset_to_initial(self) -> None: ...

    def stop(self, arms: Collection[ArmName] = ("left", "right")) -> None: ...

    def close(self) -> None: ...