"""
G1 Hardware Controller — Interface between Neon VLA and the physical robot.

Communicates with the Unitree G1 via the Unitree SDK.
Handles joint position commands, locomotion velocity, camera streaming.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple

import numpy as np

from neon.data.action_space import ControlMode, G1ActionSpace

logger = logging.getLogger(__name__)


@dataclass
class G1Config:
    """Configuration for the G1 hardware interface."""

    # Connection
    robot_ip: str = "192.168.123.10"  # Default G1 IP
    sdk_port: int = 8080
    # Control
    control_frequency: float = 50.0  # Hz (20ms per step)
    mode: str = "arms_only"  # Control mode
    # Safety
    max_joint_velocity: float = 2.0  # rad/s — safety limit
    max_linear_velocity: float = 0.5  # m/s
    max_angular_velocity: float = 0.5  # rad/s
    enable_safety_checks: bool = True
    # Camera
    camera_id: int = 0  # Which camera to use (0=head, 1=chest)
    camera_fps: int = 30
    image_size: Tuple[int, int] = (640, 480)


class G1Controller:
    """Hardware controller for the Unitree G1 humanoid.

    Provides:
    - Joint position control (arms, torso, head, legs)
    - Locomotion velocity commands
    - Camera frame capture
    - Safety limits enforcement

    Usage:
        controller = G1Controller(G1Config())
        controller.connect()

        # Control loop
        while running:
            image = controller.get_camera_frame()
            joint_states = controller.get_joint_states()

            # ... run VLA model ...
            actions = model.predict(image, instruction, joint_states)

            controller.send_actions(actions.raw_actions[0])  # First step of chunk

        controller.disconnect()
    """

    def __init__(self, config: G1Config):
        self.config = config
        self.action_space = G1ActionSpace(
            mode=ControlMode(config.mode),
            include_locomotion=True,
        )
        self._connected = False
        self._robot = None
        self._camera = None

    def connect(self):
        """Connect to the G1 robot."""
        logger.info(f"Connecting to G1 at {self.config.robot_ip}...")

        try:
            # Try Unitree SDK
            import unitree_sdk2py  # noqa: F401
            from unitree_sdk2py.core.channel import ChannelFactory
            from unitree_sdk2py.go2.sport import SportClient  # G1 uses similar API

            ChannelFactory.Init(0, self.config.robot_ip)
            self._robot = SportClient()
            self._robot.Init()
            self._connected = True
            logger.info("Connected to G1 via Unitree SDK")

        except ImportError:
            logger.warning(
                "Unitree SDK not installed. Running in simulation/mock mode. "
                "Install with: pip install unitree_sdk2py"
            )
            self._connected = False
            self._robot = None

    def disconnect(self):
        """Disconnect from the robot."""
        if self._connected:
            # Send zero velocity to stop
            self._send_locomotion_command(0.0, 0.0, 0.0)
            self._connected = False
            logger.info("Disconnected from G1")

    def is_connected(self) -> bool:
        return self._connected

    def get_camera_frame(self):
        """Capture a frame from the robot's camera.

        Returns:
            PIL Image or None
        """
        if self._camera is not None:
            try:
                import cv2
                from PIL import Image

                ret, frame = self._camera.read()
                if ret:
                    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    return Image.fromarray(frame_rgb)
            except Exception as e:
                logger.warning(f"Camera capture failed: {e}")

        # Fallback: try network camera
        try:
            import cv2
            from PIL import Image

            cap = cv2.VideoCapture(f"http://{self.config.robot_ip}:8080/camera")
            ret, frame = cap.read()
            cap.release()
            if ret:
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                return Image.fromarray(frame_rgb)
        except Exception:
            pass

        return None

    def get_joint_states(self) -> np.ndarray:
        """Get current joint positions.

        Returns:
            np.ndarray of shape (action_dim,) with current positions
        """
        if self._robot is not None:
            try:
                # Get state from Unitree SDK
                state = self._robot.GetMotorState()
                positions = np.array([s.q for s in state], dtype=np.float32)
                return positions[: self.action_space.action_dim]
            except Exception as e:
                logger.warning(f"Failed to read joint states: {e}")

        # Return default position if no robot
        return self.action_space.default_position()

    def send_actions(self, actions: np.ndarray):
        """Send action commands to the robot.

        Args:
            actions: (action_dim,) array of joint positions + locomotion velocity
                Already denormalized to physical units.
        """
        if self.config.enable_safety_checks:
            actions = self._apply_safety_limits(actions)

        # Split into joint positions and locomotion
        split = self.action_space.split_action(actions)

        # Send joint position commands
        joint_groups = ["left_arm", "right_arm", "torso", "head"]
        if self.action_space.mode == ControlMode.WHOLE_BODY:
            joint_groups.extend(["left_leg", "right_leg"])

        for group in joint_groups:
            if group in split:
                self._send_joint_command(group, split[group])

        # Send locomotion velocity
        if "locomotion" in split:
            vx, vy, omega = split["locomotion"]
            self._send_locomotion_command(vx, vy, omega)

    def _send_joint_command(self, group: str, positions: np.ndarray):
        """Send joint position command for a specific group."""
        if self._robot is None:
            return

        try:
            # Map group to Unitree motor IDs
            # This mapping depends on the specific SDK version
            indices = self.action_space.get_group_indices(group)
            for i, idx in enumerate(indices):
                if i < len(positions):
                    # Send position command via SDK
                    # self._robot.SetMotorPosition(motor_id=idx, position=positions[i])
                    pass  # Placeholder until SDK is integrated
        except Exception as e:
            logger.warning(f"Failed to send {group} command: {e}")

    def _send_locomotion_command(self, vx: float, vy: float, omega: float):
        """Send locomotion velocity command."""
        if self._robot is None:
            return

        try:
            self._robot.Move(vx, vy, omega)
        except Exception as e:
            logger.warning(f"Failed to send locomotion command: {e}")

    def _apply_safety_limits(self, actions: np.ndarray) -> np.ndarray:
        """Apply safety limits to prevent dangerous motions."""
        actions = actions.copy()

        # Clip joint velocities (approximate from position commands)
        joints = self.action_space.active_joints
        n_joints = len(joints)
        for i, joint in enumerate(joints):
            if i < len(actions):
                actions[i] = np.clip(actions[i], joint.lower_limit, joint.upper_limit)

        # Clip locomotion velocity
        if self.action_space.include_locomotion and len(actions) > n_joints:
            actions[n_joints] = np.clip(actions[n_joints], -self.config.max_linear_velocity, self.config.max_linear_velocity)
            actions[n_joints + 1] = np.clip(actions[n_joints + 1], -self.config.max_linear_velocity, self.config.max_linear_velocity)
            actions[n_joints + 2] = np.clip(actions[n_joints + 2], -self.config.max_angular_velocity, self.config.max_angular_velocity)

        return actions

    def run_control_loop(
        self,
        model,
        instruction: str = "",
        max_steps: int = 1000,
        action_chunk_index: int = 0,
    ):
        """Run a closed-loop control session.

        Args:
            model: NeonVLA model or inference server
            instruction: Language instruction
            max_steps: Maximum control steps
            action_chunk_index: Which step of the action chunk to execute (0 = first)
        """
        logger.info(f"Starting control loop: '{instruction}' (max {max_steps} steps)")

        dt = 1.0 / self.config.control_frequency
        step = 0

        try:
            while step < max_steps:
                t0 = time.time()

                # Observe
                image = self.get_camera_frame()
                joint_states = self.get_joint_states()

                # Predict
                output = model.predict(
                    image=image,
                    instruction=instruction,
                    proprioception=joint_states,
                )

                # Execute first action in chunk
                if output.raw_actions is not None and len(output.raw_actions) > action_chunk_index:
                    self.send_actions(output.raw_actions[action_chunk_index])

                step += 1

                # Control rate
                elapsed = time.time() - t0
                sleep_time = max(0, dt - elapsed)
                if sleep_time > 0:
                    time.sleep(sleep_time)

                if step % 50 == 0:
                    logger.info(f"Step {step}/{max_steps}, freq: {1.0/(elapsed+1e-6):.1f} Hz")

        except KeyboardInterrupt:
            logger.info("Control loop interrupted")
        finally:
            # Safety: stop robot
            self._send_locomotion_command(0.0, 0.0, 0.0)
            logger.info(f"Control loop ended after {step} steps")
