"""Teleop manager: executes high-rate neck (and future teleop) actions from external teleop software."""

from __future__ import annotations

import asyncio
import math
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple

from .motors.robstride import RobstrideOperationCommand
from .motors.motors_manager import MotorsManager
from feathersdk.utils.logger import info, warning, error

if TYPE_CHECKING:
    from .neck_platform import NeckPlatform
    from .torso_platform import TorsoPlatform
    from .navigation_platform import VelocityNavigationPlanner
    from .manipulation_platform import ManipulationPlatform

# Default (kp, kd) when a motor is not in ARM_OPERATION_GAINS
DEFAULT_ARM_OPERATION_KP = 30.0
DEFAULT_ARM_OPERATION_KD = 5.0

# motor_name -> (kp, kd) for arm operation commands. Edit here if config loading is not used.
ARM_OPERATION_GAINS: Dict[str, Tuple[float, float]] = {
    "SpR": (500.0, 5.0),
    "SrR": (500.0, 5.0),
    "SwR": (30.0, 5.0),
    "EpR": (300.0, 5.0),
    "WwR": (30.0, 5.0),
    "WpR": (30.0, 5.0),
    "WrR": (30.0, 5.0),
    "SpL": (500.0, 5.0),
    "SrL": (500.0, 5.0),
    "SwL": (30.0, 5.0),
    "EpL": (300.0, 5.0),
    "WwL": (30.0, 5.0),
    "WpL": (30.0, 5.0),
    "WrL": (30.0, 5.0),
}

# If True, on_arm_action uses execute_arm_go_to (trajectory); if False, uses execute_operation_commands.
USE_GO_TO = False


@dataclass
class NeckAction:
    """Neck command - arrives at ~120Hz (future feature).

    Teleop software sends this; the teleop manager forwards it to the neck platform.
    """

    joint_names: List[str]   # e.g. ["neck_pan", "neck_tilt"]
    positions: List[float]   # radians
    velocities: List[float]  # rad/s
    efforts: List[float]     # Target efforts in Nm (will be zeros)
    timestamp: float


@dataclass
class ElevatorAction:
    """Elevator command - arrives at ~30Hz continuously.

    Teleop software sends this; the teleop manager forwards it to the torso platform.
    """

    velocity: float   # Normalized z velocity -1.0 to 1.0 (0=stop, positive=up, negative=down; can be binary)
    timestamp: float


@dataclass
class LocomotionAction:
    """Locomotion command - arrives at ~30Hz continuously (even when zero).

    Teleop software sends this; the teleop manager forwards it to the nav planner.
    """

    linear_x: float   # Forward/backward, normalized -1.0 to 1.0
    linear_y: float  # Left/right, normalized -1.0 to 1.0
    angular_z: float  # Rotation, normalized -1.0 to 1.0
    timestamp: float


@dataclass
class ArmAction:
    """Arm command - arrives at ~120Hz continuously.

    Teleop software sends this; the teleop manager forwards it to the arm motors
    via operation commands (target angle, velocity, effort, kp/kd).
    """

    joint_names: List[str]   # e.g. ["SpL", "SrL", "SwL", ...]
    positions: List[float]   # Target positions in radians
    velocities: List[float]  # Target velocities in rad/s (often zeros)
    efforts: List[float]    # Target efforts in Nm (often zeros)
    timestamp: float        # When command was received (seconds since epoch)


class TeleopManager:
    """
    All execute_* methods are non-blocking: they set targets or queue commands and return
    immediately. Actual motion is performed by each platform's step loop or background thread.
    """

    def __init__(
        self,
        neck: Optional["NeckPlatform"] = None,
        torso: Optional["TorsoPlatform"] = None,
        nav_planner: Optional["VelocityNavigationPlanner"] = None,
        manipulation_platform: Optional["ManipulationPlatform"] = None,
        motors_manager: Optional[MotorsManager] = None,
        config: Optional[dict] = None,
    ):
        """Initialize with the neck, torso, nav planner, and/or motors manager to control.
        config: optional dict for torso_max_velocity_mm_s, elevator_deadzone, locomotion_deadzone,
            locomotion_max_linear_m_s, locomotion_max_rotation_deg_s. Arm gains use module-level
            ARM_OPERATION_GAINS and DEFAULT_ARM_OPERATION_KP/KD.
        """
        cfg = config if isinstance(config, dict) else {}
        self._neck = neck
        self._torso = torso
        self._nav_planner = nav_planner
        self._manipulation_platform = manipulation_platform
        self._motors_manager = motors_manager
        self._default_arm_operation_kp = DEFAULT_ARM_OPERATION_KP
        self._default_arm_operation_kd = DEFAULT_ARM_OPERATION_KD
        self.arm_operation_gains = dict(ARM_OPERATION_GAINS)
        self._torso_max_velocity_mm_s = float(cfg.get("torso_max_velocity_mm_s", 50.0))
        self._elevator_deadzone = float(cfg.get("elevator_deadzone", 0.2))
        self._locomotion_deadzone = float(cfg.get("locomotion_deadzone", 0.2))
        self._locomotion_max_linear_m_s = float(cfg.get("locomotion_max_linear_m_s", 0.3))
        self._locomotion_max_rotation_deg_s = float(cfg.get("locomotion_max_rotation_deg_s", 100.0))
        # Cache last-sent values to avoid redundant downstream calls.
        self._last_torso_velocity: Optional[float] = None
        self._last_neck: Optional[Tuple[Optional[float], Optional[float]]] = None
        self._last_locomotion: Optional[Tuple[str, float]] = None  # ('stop'|'linear_x'|'linear_y'|'angular_z', value)
        self._last_arm: Optional[Tuple[Tuple[str, ...], Tuple[float, ...], Tuple[float, ...], Tuple[float, ...]]] = None
        self._last_open_left: Optional[float] = None
        self._last_open_right: Optional[float] = None
        self._last_grip_left: Optional[float] = None
        self._last_grip_right: Optional[float] = None

        # if self._manipulation_platform is not None:
        #     warning(f"Enabling arms")
        #     self._manipulation_platform.right_arm.enable()
        #     self._manipulation_platform.left_arm.enable()
        #     kp_by_motor = {m: ARM_OPERATION_GAINS[m][0] for m in ARM_OPERATION_GAINS}
        #     kd_by_motor = {m: ARM_OPERATION_GAINS[m][1] for m in ARM_OPERATION_GAINS}
        #     self._manipulation_platform.right_arm.freeze_current_arm_state(kp_by_motor, kd_by_motor)
        #     self._manipulation_platform.left_arm.freeze_current_arm_state(kp_by_motor, kd_by_motor)
        #     time.sleep(1)
        #     self._manipulation_platform.right_arm.home()
        #     self._manipulation_platform.left_arm.home()
       

    def on_open_trigger_change(self, is_right: bool = True, percentage: float = 0.0) -> None:
        """Handle open-trigger change from teleop. Prints only when value changed."""
        pct = percentage
        if is_right:
            if self._last_open_right == pct:
                return
            self._last_open_right = pct
        else:
            if self._last_open_left == pct:
                return
            self._last_open_left = pct
        side = "right" if is_right else "left"
        print(f"open_trigger {side}: {pct}")
        warning(f"open_trigger {side}: {pct}")

    def on_grip_trigger_change(self, is_right: bool = True, percentage: float = 0.0) -> None:
        """Handle grip-trigger change from teleop. Prints only when value changed."""
        pct = percentage
        if is_right:
            if self._last_grip_right == pct:
                return
            self._last_grip_right = pct
        else:
            if self._last_grip_left == pct:
                return
            self._last_grip_left = pct
        side = "right" if is_right else "left"
        print(f"grip_trigger {side}: {pct}")

    def on_y_press(self) -> None:
        """Handle Y button press from teleop. Override or hook as needed."""
        if self._manipulation_platform is not None:
            self._motors_manager.trigger_compliance_mode("right_arm", "SpR")
            self._motors_manager.trigger_compliance_mode("left_arm", "SpL")
        if self._torso is not None:
            self._torso.update_velocity(0)
        if self._nav_planner is not None:
            self._nav_planner.set_speed(0)
        print("Y pressed")

    def start_dataset(self) -> None:
        """Start dataset recording. Override or hook as needed."""
        print("start_dataset")

    def stop_dataset(self) -> None:
        """Stop dataset recording. Override or hook as needed."""
        print("stop_dataset")

    def execute_neck_recommended_action(self, action: NeckAction) -> None:
        """Apply a single NeckAction to the neck platform. Safe to call at ~120Hz.
        """
        warning(f"{action.timestamp} execute_neck_recommended_action: {action}")
        neck = self._neck
        if neck is None:
            return

        pitch_rad: Optional[float] = None
        yaw_rad: Optional[float] = None
        names = action.joint_names
        positions = action.positions
        n = len(names)
        if n != len(positions):
            return

        for i in range(n):
            name = names[i]
            pos = positions[i]
            if name == neck.PITCH_MOTOR_NAME:
                pitch_rad = pos
            elif name == neck.YAW_MOTOR_NAME:
                yaw_rad = pos

        if pitch_rad is None and yaw_rad is None:
            return

        pitch_deg = math.degrees(pitch_rad) if pitch_rad is not None else None
        yaw_deg = math.degrees(yaw_rad) if yaw_rad is not None else None
        key = (pitch_deg, yaw_deg)
        if self._last_neck == key:
            return
        self._last_neck = key
        neck.go_to(pitch_in_degrees=pitch_deg, yaw_in_degrees=yaw_deg)

    def execute_torso_recommended_action(self, action: ElevatorAction) -> None:
        """Apply a single ElevatorAction to the torso platform. Safe to call at ~30Hz.
        """
        warning(f"{action.timestamp} execute_torso_recommended_action: {action}")
        torso = self._torso
        if torso is None:
            return
        v = max(-1.0, min(1.0, float(action.velocity)))
        # Deadzone: treat stick drift as stop so torso does not creep
        if -self._elevator_deadzone <= v <= self._elevator_deadzone:
            v = 0.0
        mm_per_second = v * self._torso_max_velocity_mm_s
        if self._last_torso_velocity == mm_per_second:
            return
        self._last_torso_velocity = mm_per_second
        try:
            torso.update_velocity(mm_per_second)
        except Exception:
            pass  # e.g. not healthy, or at limits; torso already no-ops at limits
        

    def execute_arm_recommended_action(self, action: ArmAction) -> None:
        """Apply a single ArmAction to arm motors via operation commands. Safe to call at ~120Hz.
        """
        if USE_GO_TO:
            self.execute_arm_go_to(action) # 120Hz
        else:
            self.execute_operation_commands(action)

    def execute_operation_commands(self, action: ArmAction) -> None:
        mm = self._motors_manager
        if mm is None:
            return
        names = action.joint_names
        positions = action.positions
        velocities = action.velocities
        efforts = action.efforts

        n = len(names)
        if n != len(positions):
            error(f"{action.timestamp} execute_arm_recommended_action: len(positions) != len(names)")
            return
        key = (tuple(names), tuple(positions), tuple(velocities), tuple(efforts))
        if self._last_arm == key:
            return
        self._last_arm = key
        default_gains = (self._default_arm_operation_kp, self._default_arm_operation_kd)
        gains = self.arm_operation_gains
        commands: List[RobstrideOperationCommand] = []
        for i in range(n):
            motor_name = names[i]
            try:
                motor = mm.get_motor(motor_name)
            except Exception as e:
                error(f"error getting motor {motor_name}: {e}")
            kp, kd = gains.get(motor_name, default_gains)

            warning(f"SENDING operation command to {motor_name}: {motor.convert_calibrated_angle_to_angle(positions[i])} kp: {kp} kd: {kd}")
            cmd = RobstrideOperationCommand(
                motor,
                0,
                motor.convert_calibrated_angle_to_angle(positions[i]),
                velocities[i],
                kp,
                kd,
            )
            commands.append(cmd)
        if commands:
            try:
                mm.send_operation_commands(commands)
            except Exception as e:
                error(f"{action.timestamp} execute_arm_recommended_action: {e}")
                pass

    def execute_arm_go_to(self, action: ArmAction, total_time: float = 1.0) -> None:
        """Apply ArmAction to arms using arm.go_to (trajectory). Non-blocking.

        Builds motor name -> target position from action.joint_names and action.positions
        (radians / calibrated angle), then calls left_arm.go_to and right_arm.go_to.
        """
        warning(f"execute_arm_go_to: {action}")
        platform = self._manipulation_platform
        if platform is None:
            return
        names = action.joint_names
        positions = action.positions
        if len(names) != len(positions):
            error(f"execute_arm_go_to: len(joint_names) != len(positions)")
            return
        motors_targets = dict(zip(names, positions))
        # Use each arm's motor names as source of truth (arm_system.motors is Dict[str, Motor]).
        right_targets = {m: motors_targets[m] for m in platform.right_arm.motors if m in motors_targets}
        left_targets = {m: motors_targets[m] for m in platform.left_arm.motors if m in motors_targets}
        if not right_targets and not left_targets:
            warning(
                f"execute_arm_go_to: no targets for any arm. "
                f"right_arm.motors keys={list(platform.right_arm.motors.keys())!r}, "
                f"left_arm.motors keys={list(platform.left_arm.motors.keys())!r}, "
                f"action joint_names={names!r}"
            )
        if right_targets:
            try:
                warning(f"execute_arm_go_to right_arm: {right_targets} {total_time}")
                platform.right_arm.go_to(right_targets, total_time)
            except Exception as e:
                error(f"execute_arm_go_to right_arm: {e}")
        if left_targets:
            try:
                warning(f"execute_arm_go_to left_arm: {left_targets} {total_time}")
                platform.left_arm.go_to(left_targets, total_time)
            except Exception as e:
                error(f"execute_arm_go_to left_arm: {e}")

    def execute_base_recommended_action(self, action: LocomotionAction) -> None:
        """Apply a single LocomotionAction to the nav planner. Safe to call at ~30Hz.
        """
        warning(f"{action.timestamp} execute_base_recommended_action: {action}")
        nav = self._nav_planner
        if nav is None:
            return
        dx = max(-1.0, min(1.0, float(action.linear_x)))
        dy = max(-1.0, min(1.0, float(action.linear_y)))
        dz = max(-1.0, min(1.0, float(action.angular_z)))
        # Apply deadzone: treat as zero inside [-locomotion_deadzone, locomotion_deadzone]
        def above_deadzone(v: float) -> float:
            if -self._locomotion_deadzone <= v <= self._locomotion_deadzone:
                return 0.0
            return v
        ax = above_deadzone(dx)
        ay = above_deadzone(dy)
        az = above_deadzone(dz)
        mag_x = abs(ax)
        mag_y = abs(ay)
        mag_z = abs(az)
        if mag_x == 0 and mag_y == 0 and mag_z == 0:
            cmd = ("stop", 0.0)
        elif mag_x >= mag_y and mag_x >= mag_z:
            cmd = ("linear_x", ax * self._locomotion_max_linear_m_s)
        elif mag_y >= mag_z:
            cmd = ("linear_y", ay * self._locomotion_max_linear_m_s)
        else:
            cmd = ("angular_z", -az * self._locomotion_max_rotation_deg_s)
        if self._last_locomotion == cmd:
            return
        self._last_locomotion = cmd
        try:
            if cmd[0] == "stop":
                nav.set_speed(0)
            elif cmd[0] == "linear_x":
                nav.go_to(0, cmd[1])
            elif cmd[0] == "linear_y":
                nav.go_to(90, cmd[1])
            else:
                nav.rotate_in_place(cmd[1])
        except Exception:
            pass
    
    
    def arm_observation(self) -> Dict[str, dict]:
        """Read current state of arms. Returns motor_name -> {temp, torque, angle, velocity}."""
        platform = self._manipulation_platform
        if platform is None:
            return {}
        out: Dict[str, dict] = {}
        out.update(platform.right_arm.get_state())
        out.update(platform.left_arm.get_state())
        return out

    def handle_dataset_start(self, dataset_id: str) -> None:
        print(f"Dataset {dataset_id} started")
        warning(f"Dataset {dataset_id} started")

    def handle_dataset_stop(self, dataset_id: str) -> None:
        print(f"Dataset {dataset_id} stopped")
        warning(f"Dataset {dataset_id} stopped")

    def on_y_pressed(self) -> None:
        print("Y pressed, emergency stop")
        warning("Y pressed, emergency stop")