#!/usr/bin/env python

# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Docs:
# hebi: https://docs.hebi.us/tools.html#mobile-io
# teleop: https://github.com/SpesRobotics/teleop

import logging
import threading
import time
from typing import TYPE_CHECKING, Any

import numpy as np

from lerobot.lerobot_types import RobotAction
from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
from lerobot.utils.errors import DeviceNotConnectedError
from lerobot.utils.import_utils import _hebi_available, _teleop_available, require_package
from lerobot.utils.rotation import Rotation

if TYPE_CHECKING or _hebi_available:
    import hebi
else:
    hebi = None

if TYPE_CHECKING or _teleop_available:
    from teleop import Teleop
else:
    Teleop = None

if TYPE_CHECKING:
    from hebi._internal.group import Group as HebiGroup

from ..teleoperator import Teleoperator
from .config_phone import PhoneConfig, PhoneOS

logger = logging.getLogger(__name__)


class BasePhone:
    _enabled: bool = False
    _calib_pos: np.ndarray | None = None
    _calib_rot_inv: Rotation | None = None

    def _reapply_position_calibration(self, pos: np.ndarray) -> None:
        self._calib_pos = pos.copy()

    @property
    def is_calibrated(self) -> bool:
        return (self._calib_pos is not None) and (self._calib_rot_inv is not None)

    @property
    def action_features(self) -> dict[str, type]:
        return {
            "phone.pos": np.ndarray,  # shape (3,)
            "phone.rot": Rotation,  # scipy.spatial.transform.Rotation
            "phone.raw_inputs": dict,  # analogs/buttons or webXR meta
            "phone.enabled": bool,
        }

    @property
    def feedback_features(self) -> dict[str, type]:
        # No haptic or other feedback implemented yet
        return {}

    def configure(self) -> None:
        # No additional configuration required for phone teleop
        pass

    def send_feedback(self, feedback: dict[str, float]) -> None:
        # We could add haptic feedback (vibrations) here, but it's not implemented yet
        raise NotImplementedError

    def _apply_calibration(self, raw_pos: np.ndarray, raw_rot: Rotation) -> tuple[np.ndarray, Rotation]:
        """Express a raw phone pose in the reference frame captured by `calibrate()`."""
        if self._calib_pos is None or self._calib_rot_inv is None:
            raise RuntimeError(f"{self} is not calibrated. Run `.calibrate()` first.")
        pos_cal = self._calib_rot_inv.apply(raw_pos - self._calib_pos)
        rot_cal = self._calib_rot_inv * raw_rot
        return pos_cal, rot_cal


class IOSPhone(BasePhone, Teleoperator):
    name = "ios_phone"

    def __init__(self, config: PhoneConfig):
        require_package("hebi-py", extra="phone", import_name="hebi")
        require_package("teleop", extra="phone")
        super().__init__(config)
        self.config = config
        self._group: HebiGroup | None = None

    @property
    def is_connected(self) -> bool:
        return self._group is not None

    @check_if_already_connected
    def connect(self, calibrate: bool = True) -> None:
        logger.info("Connecting to IPhone, make sure to open the HEBI Mobile I/O app.")
        lookup = hebi.Lookup()
        time.sleep(2.0)
        group = lookup.get_group_from_names(["HEBI"], ["mobileIO"])
        if group is None:
            raise RuntimeError("Mobile I/O not found — check name/family settings in the app.")
        self._group = group
        logger.info(f"{self} connected to HEBI group with {group.size} module(s).")

        if calibrate:
            self.calibrate()

    def calibrate(self) -> None:
        print(
            "Hold the phone so that: top edge points forward in same direction as the robot (robot +x) and screen points up (robot +z)"
        )
        print("Press and hold B1 in the HEBI Mobile I/O app to capture this pose...\n")
        position, rotation = self._wait_for_capture_trigger()
        self._calib_pos = position.copy()
        self._calib_rot_inv = rotation.inv()
        self._enabled = False
        print("Calibration done\n")

    def _wait_for_capture_trigger(self) -> tuple[np.ndarray, Rotation]:
        """
        Blocks execution until the calibration trigger is detected from the iOS device.

        This method enters a loop, continuously reading the phone's state. It waits for the user to press
        and hold the 'B1' button in the HEBI Mobile I/O app. Once B1 is pressed, the loop breaks and
        returns the phone's pose at that exact moment.

        Returns:
            A tuple containing the position (np.ndarray) and rotation (Rotation) of the phone at the
            moment the trigger was activated.
        """
        while True:
            pose = self._read_current_pose()
            if pose is None:
                time.sleep(0.01)
                continue
            position, rotation, fb_pose = pose

            io = getattr(fb_pose, "io", None)
            button_b = getattr(io, "b", None) if io is not None else None
            button_b1_pressed = False
            if button_b is not None:
                button_b1_pressed = bool(button_b.get_int(1))
            if button_b1_pressed:
                return position, rotation

            time.sleep(0.01)

    def _read_current_pose(self) -> tuple[np.ndarray, Rotation, object] | None:
        """
        Reads the instantaneous 6-DoF pose from the connected iOS device via the HEBI SDK.

        This method fetches the latest feedback packet from the HEBI group, extracts the ARKit
        position and orientation, and converts them into a standard format. It also applies a
        configured camera offset to adjust the pose from the camera's frame to the phone's
        physical frame.

        Returns:
            None if no valid pose could be read, otherwise a tuple containing:
            - The 3D position as a NumPy array.
            - The orientation as a `Rotation` object.
            - The raw HEBI feedback object for accessing other data like button presses.
        """
        if self._group is None:
            raise DeviceNotConnectedError(f"{self} is not connected. Run `.connect()` first.")
        fbk = self._group.get_next_feedback()
        pose = fbk[0]
        ar_pos = getattr(pose, "ar_position", None)
        ar_quat = getattr(pose, "ar_orientation", None)
        if ar_pos is None or ar_quat is None:
            return None
        # HEBI provides orientation in w, x, y, z format.
        # Scipy's Rotation expects x, y, z, w.
        quat_xyzw = np.concatenate((ar_quat[1:], [ar_quat[0]]))  # wxyz to xyzw
        # ARKit can emit zero/NaN quaternions before tracking is ready or on a
        # dropped packet. Rotation.from_quat now rejects those; degrade the same
        # way as a missing pose so teleop stays alive mid-session.
        try:
            rot = Rotation.from_quat(quat_xyzw)
        except ValueError:
            return None
        pos = ar_pos - rot.apply(self.config.camera_offset)
        return pos, rot, pose

    @check_if_not_connected
    def get_action(self) -> RobotAction:
        pose = self._read_current_pose()
        if pose is None or not self.is_calibrated:
            return {}
        raw_position, raw_rotation, fb_pose = pose

        # Collect raw inputs (B1 / analogs on iOS, move/scale on Android)
        raw_inputs: dict[str, float | int | bool] = {}
        io = getattr(fb_pose, "io", None)
        if io is not None:
            bank_a, bank_b = io.a, io.b
            if bank_a:
                for ch in range(1, 9):
                    if bank_a.has_float(ch):
                        raw_inputs[f"a{ch}"] = float(bank_a.get_float(ch))
            if bank_b:
                for ch in range(1, 9):
                    if bank_b.has_int(ch):
                        raw_inputs[f"b{ch}"] = int(bank_b.get_int(ch))
                    elif hasattr(bank_b, "has_bool") and bank_b.has_bool(ch):
                        raw_inputs[f"b{ch}"] = int(bank_b.get_bool(ch))

        enable = bool(raw_inputs.get("b1", 0))

        # Rising edge then re-capture calibration immediately from current raw pose
        if enable and not self._enabled:
            self._reapply_position_calibration(raw_position)

        pos_cal, rot_cal = self._apply_calibration(raw_position, raw_rotation)

        self._enabled = enable

        return {
            "phone.pos": pos_cal,
            "phone.rot": rot_cal,
            "phone.raw_inputs": raw_inputs,
            "phone.enabled": self._enabled,
        }

    @check_if_not_connected
    def disconnect(self) -> None:
        self._group = None


class AndroidPhone(BasePhone, Teleoperator):
    name = "android_phone"

    def __init__(self, config: PhoneConfig):
        require_package("hebi-py", extra="phone", import_name="hebi")
        require_package("teleop", extra="phone")
        super().__init__(config)
        self.config = config
        self._teleop: Teleop | None = None
        self._teleop_thread: threading.Thread | None = None
        self._latest_pose: np.ndarray | None = None
        self._latest_message: dict[str, Any] | None = None
        self._android_lock = threading.Lock()

    @property
    def is_connected(self) -> bool:
        return self._teleop is not None

    @check_if_already_connected
    def connect(self, calibrate: bool = True) -> None:
        logger.info("Starting teleop stream for Android...")
        self._teleop = Teleop()
        self._teleop.subscribe(self._android_callback)
        self._teleop_thread = threading.Thread(target=self._teleop.run, daemon=True)
        self._teleop_thread.start()
        logger.info(f"{self} connected, teleop stream started.")

        if calibrate:
            self.calibrate()

    def calibrate(self) -> None:
        print(
            "Hold the phone so that: top edge points forward in same direction as the robot (robot +x) and screen points up (robot +z)"
        )
        print("Touch and move on the WebXR page to capture this pose...\n")

        pos, rot = self._wait_for_capture_trigger()
        self._calib_pos = pos.copy()
        self._calib_rot_inv = rot.inv()
        self._enabled = False
        print("Calibration done\n")

    def _wait_for_capture_trigger(self) -> tuple[np.ndarray, Rotation]:
        """
        Blocks execution until the calibration trigger is detected from the Android device.

        This method enters a loop, continuously checking the latest message received from the WebXR
        session. It waits for the user to touch and move their finger on the screen, which generates
        a `move` event. Once this event is detected, the loop breaks and returns the phone's current
        pose.

        Returns:
            A tuple containing the position (np.ndarray) and rotation (Rotation) of the phone at the
            moment the trigger was activated.
        """
        while True:
            with self._android_lock:
                msg = self._latest_message or {}

            if bool(msg.get("move", False)):
                pose = self._read_current_pose()
                if pose is not None:
                    pos, rot, _raw_pose = pose
                    return pos, rot

            time.sleep(0.01)

    def _read_current_pose(self) -> tuple[np.ndarray, Rotation, np.ndarray] | None:
        """
        Reads the latest 6-DoF pose received from the Android device's WebXR session.

        This method accesses the most recent pose data stored by the `_android_callback`. It uses a
        thread lock to safely read the shared `_latest_pose` variable. The pose, a 4x4 matrix, is
        then decomposed into position and rotation, and the configured camera offset is applied.

        Returns:
            None if no pose has been received yet, otherwise a tuple containing:
            - The 3D position as a NumPy array.
            - The orientation as a `Rotation` object.
            - The raw 4x4 pose matrix as received from the teleop stream.
        """
        with self._android_lock:
            if self._latest_pose is None:
                return None
            p = self._latest_pose.copy()
            pose = self._latest_pose
        rot = Rotation.from_matrix(p[:3, :3])
        pos = p[:3, 3] - rot.apply(self.config.camera_offset)
        return pos, rot, pose

    def _android_callback(self, pose: np.ndarray, message: dict[str, Any]) -> None:
        """
        Callback function to handle incoming data from the Android teleop stream.

        This method is executed by the `teleop` package's subscriber thread whenever a new
        pose and message are received from the WebXR session on the Android phone. It updates
        the internal state (`_latest_pose` and `_latest_message`) with the new data.
        A thread lock is used to ensure that these shared variables are updated atomically,
        preventing race conditions with the main thread that reads them.

        Args:
            pose: A 4x4 NumPy array representing the phone's transformation matrix.
            message: A dictionary containing additional data, such as button presses or touch events.
        """
        with self._android_lock:
            self._latest_pose = pose
            self._latest_message = message

    @check_if_not_connected
    def get_action(self) -> RobotAction:
        pose = self._read_current_pose()
        if pose is None or not self.is_calibrated:
            return {}
        raw_pos, raw_rot, _raw_pose = pose

        # Collect raw inputs (B1 / analogs on iOS, move/scale on Android)
        raw_inputs: dict[str, float | int | bool] = {}
        msg = self._latest_message or {}
        raw_inputs["move"] = bool(msg.get("move", False))
        raw_inputs["scale"] = float(msg.get("scale", 1.0))
        raw_inputs["reservedButtonA"] = bool(msg.get("reservedButtonA", False))
        raw_inputs["reservedButtonB"] = bool(msg.get("reservedButtonB", False))

        enable = bool(raw_inputs.get("move", False))

        # Rising edge then re-capture calibration immediately from current raw pose
        if enable and not self._enabled:
            self._reapply_position_calibration(raw_pos)

        pos_cal, rot_cal = self._apply_calibration(raw_pos, raw_rot)

        self._enabled = enable

        return {
            "phone.pos": pos_cal,
            "phone.rot": rot_cal,
            "phone.raw_inputs": raw_inputs,
            "phone.enabled": self._enabled,
        }

    @check_if_not_connected
    def disconnect(self) -> None:
        self._teleop = None
        if self._teleop_thread and self._teleop_thread.is_alive():
            self._teleop_thread.join(timeout=1.0)
            self._teleop_thread = None
            self._latest_pose = None


class Phone(Teleoperator):
    """
    Phone-based teleoperator using ARKit (iOS via HEBI Mobile I/O App) or the teleop Python package (Android via WebXR API).
    For HEBI Mobile I/O we also expose 8 analog (a1-a8) and 8 digital (b1-b8) inputs.

    Press and hold **B1** to enable teleoperation. While enabled, the first B1 press
    captures a reference pose and rotation, when disabled and pressed again the position is reapplied.
    """

    config_class = PhoneConfig
    name = "phone"

    def __init__(self, config: PhoneConfig):
        super().__init__(config)
        self.config = config

        self._phone_impl: Teleoperator

        if self.config.phone_os == PhoneOS.IOS:
            self._phone_impl = IOSPhone(config)
        elif self.config.phone_os == PhoneOS.ANDROID:
            self._phone_impl = AndroidPhone(config)
        else:
            raise ValueError(f"Invalid config phone_os: {self.config.phone_os}")

    @property
    def is_connected(self) -> bool:
        return self._phone_impl.is_connected

    def connect(self, calibrate: bool = True) -> None:
        return self._phone_impl.connect(calibrate)

    def calibrate(self) -> None:
        return self._phone_impl.calibrate()

    @property
    def is_calibrated(self) -> bool:
        return self._phone_impl.is_calibrated

    @property
    def action_features(self) -> dict[str, type]:
        return self._phone_impl.action_features

    @property
    def feedback_features(self) -> dict[str, type]:
        return self._phone_impl.feedback_features

    def configure(self) -> None:
        return self._phone_impl.configure()

    def get_action(self) -> RobotAction:
        return self._phone_impl.get_action()

    def send_feedback(self, feedback: dict[str, float]) -> None:
        return self._phone_impl.send_feedback(feedback)

    def disconnect(self) -> None:
        return self._phone_impl.disconnect()
