"""R1 26-DoF motor → name / mesh mapping for the 3D temperature dashboard.

Keys are the rt/lowstate motor_state SLOT index (== the "motor_id" the arm
controller emits). The R1's 26 joints do NOT occupy contiguous slots — the
authoritative layout is unitree_sdk2 example/r1 `joint_idx_in_idl` +
`R1JointIndex`: legs 0-11, waist_roll=12/waist_yaw=13, left arm 15-19, right arm
22-26, head_pitch=29/head_yaw=30 (slots 14,20,21,27,28 unused). Link names match
static/temp3d/r1/R1.urdf. `build_payload()` turns the arm controller's lowstate
snapshot into the 'motor_update' payload the three.js front-end expects.
"""

from __future__ import annotations

from typing import Any, Optional

# The 26 rt/lowstate slots the R1 populates, in body order — the arm controller
# samples exactly these (so Total Motors == 26 and every entry maps to a mesh).
R1_MOTOR_SLOTS: list[int] = [
    0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11,   # legs (L then R)
    12, 13,                                  # waist roll, yaw
    15, 16, 17, 18, 19,                      # left arm: shoulder p/r/y, elbow, wrist_roll
    22, 23, 24, 25, 26,                      # right arm
    29, 30,                                  # head pitch, yaw
]

# rt/lowstate slot → human name
MOTOR_NAMES: dict[int, str] = {
    0: "Left Hip Pitch", 1: "Left Hip Roll", 2: "Left Hip Yaw", 3: "Left Knee",
    4: "Left Ankle Pitch", 5: "Left Ankle Roll",
    6: "Right Hip Pitch", 7: "Right Hip Roll", 8: "Right Hip Yaw", 9: "Right Knee",
    10: "Right Ankle Pitch", 11: "Right Ankle Roll",
    12: "Waist Roll", 13: "Waist Yaw",
    15: "Left Shoulder Pitch", 16: "Left Shoulder Roll", 17: "Left Shoulder Yaw",
    18: "Left Elbow", 19: "Left Wrist Roll",
    22: "Right Shoulder Pitch", 23: "Right Shoulder Roll", 24: "Right Shoulder Yaw",
    25: "Right Elbow", 26: "Right Wrist Roll",
    29: "Head Pitch", 30: "Head Yaw",
}

# rt/lowstate slot → R1 URDF link / STL mesh name
MOTOR_TO_MESH: dict[int, str] = {
    0: "left_hip_pitch_link", 1: "left_hip_roll_link", 2: "left_hip_yaw_link",
    3: "left_knee_link", 4: "left_ankle_pitch_link", 5: "left_ankle_roll_link",
    6: "right_hip_pitch_link", 7: "right_hip_roll_link", 8: "right_hip_yaw_link",
    9: "right_knee_link", 10: "right_ankle_pitch_link", 11: "right_ankle_roll_link",
    12: "waist_roll_link", 13: "waist_yaw_link",
    15: "left_shoulder_pitch_link", 16: "left_shoulder_roll_link", 17: "left_shoulder_yaw_link",
    18: "left_elbow_link", 19: "left_wrist_roll_link",
    22: "right_shoulder_pitch_link", 23: "right_shoulder_roll_link", 24: "right_shoulder_yaw_link",
    25: "right_elbow_link", 26: "right_wrist_roll_link",
    29: "head_pitch_link", 30: "head_yaw_link",
}

# Temperature thresholds (°C) — the three.js gradient maps MIN→MAX (blue→red).
TEMP_MIN = 30
TEMP_MAX = 120
TEMP_WARM_THRESHOLD = 45
TEMP_HOT_THRESHOLD = 60


def _coerce(v: Optional[int]) -> float:
    """Temperatures default to 0 when the firmware didn't report one, so the
    front-end's Math.max / .toFixed never sees null/NaN."""
    return float(v) if v is not None else 0.0


def build_payload(temps: list[dict[str, Any]],
                  positions: list[float],
                  timestamp: float) -> dict[str, Any]:
    """Build the Marcus-compatible 'motor_update' payload.

    `temps`     — arm.get_motor_temps(): [{motor_id, surface, winding}]
    `positions` — arm.get_current_q(): joint angles indexed by motor id
    """
    temperatures: list[dict[str, Any]] = []
    for t in temps or []:
        i = t.get("motor_id")
        surface = t.get("surface")
        winding = t.get("winding")
        if surface is not None and winding is not None:
            avg = (_coerce(surface) + _coerce(winding)) / 2.0
        else:
            avg = _coerce(surface if surface is not None else winding)
        entry: dict[str, Any] = {
            "motor_id": i,
            "motor_name": MOTOR_NAMES.get(i, f"Motor {i}"),
            "mesh_name": MOTOR_TO_MESH.get(i, ""),
            "surface": _coerce(surface),
            "winding": _coerce(winding),
            "temp1": _coerce(surface),
            "temp2": _coerce(winding),
            "avg": avg,
        }
        if positions and isinstance(i, int) and i < len(positions):
            entry["position"] = float(positions[i])
        temperatures.append(entry)

    pos_list: list[dict[str, Any]] = [
        {"motor_id": i, "position": float(q), "link_name": MOTOR_TO_MESH.get(i)}
        for i, q in enumerate(positions or [])
    ]
    return {"temperatures": temperatures, "positions": pos_list,
            "timestamp": timestamp}