import threading
import time
from typing import Dict, List

import rclpy
from rclpy.node import Node
from control_msgs.msg import JointTrajectoryControllerState
from std_srvs.srv import SetBool, Trigger
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
import yaml


class JointTrajectoryRecorder(Node):
    def __init__(self):
        super().__init__('joint_trajectory_recorder')

        self.declare_parameter('input_topic', '/left_joint_trajectory_controller/controller_state')
        self.declare_parameter('output_topic', '/left_joint_trajectory_controller/joint_trajectory')
        self.declare_parameter('play_speed', 1.0)
        self.declare_parameter('loop_playback', False)
        self.declare_parameter('storage_path', '/tmp/openarm_joint_record.yaml')

        self.input_topic = self.get_parameter('input_topic').value
        self.output_topic = self.get_parameter('output_topic').value

        self.record_lock = threading.Lock()
        self.play_lock = threading.Lock()

        self.is_recording = False
        self.is_playing = False
        self.stop_playback_event = threading.Event()

        self._record_start_time = None
        self.recorded_points: List[Dict] = []
        self.get_logger().info(f"input_topic: {self.input_topic}, output_topic: {self.output_topic}")

        self.sub = self.create_subscription(
            JointTrajectoryControllerState,
            self.input_topic,
            self.controller_state_callback,
            50,
        )
        self.pub = self.create_publisher(JointTrajectory, self.output_topic, 50)

        self.set_recording_srv = self.create_service(
            SetBool,
            '~/set_recording',
            self.handle_set_recording,
        )
        self.play_srv = self.create_service(
            Trigger,
            '~/play',
            self.handle_play,
        )
        self.stop_play_srv = self.create_service(
            Trigger,
            '~/stop_playback',
            self.handle_stop_playback,
        )
        self.save_srv = self.create_service(
            Trigger,
            '~/save',
            self.handle_save,
        )
        self.load_srv = self.create_service(
            Trigger,
            '~/load',
            self.handle_load,
        )
        self.clear_srv = self.create_service(
            Trigger,
            '~/clear',
            self.handle_clear,
        )

        # self.get_logger().info(
        #     'joint_trajectory_recorder started. Use services: '
        #     '~/set_recording, ~/play, ~/stop_playback, ~/save, ~/load, ~/clear'
        # )

    def controller_state_callback(self, msg: JointTrajectoryControllerState):
        if not self.is_recording:
            return

        names = list(msg.joint_names)
        positions = [float(v) for v in msg.actual.positions]
        if len(names) == 0 or len(positions) == 0:
            return

        now_s = time.time()
        with self.record_lock:
            if self._record_start_time is None:
                self._record_start_time = now_s
            rel_time = now_s - self._record_start_time
            self.recorded_points.append({
                't': rel_time,
                'name': names,
                'position': positions,
            })

    def handle_set_recording(self, request, response):
        if request.data:
            with self.record_lock:
                self.recorded_points = []
                self._record_start_time = None
            self.is_recording = True
            response.success = True
            response.message = 'Recording started. Existing trajectory cleared.'
            self.get_logger().info(response.message)
            return response

        self.is_recording = False
        with self.record_lock:
            count = len(self.recorded_points)
        response.success = True
        response.message = f'Recording stopped. Collected {count} points.'
        self.get_logger().info(response.message)
        return response

    def handle_play(self, _request, response):
        with self.play_lock:
            if self.is_playing:
                response.success = False
                response.message = 'Playback is already running.'
                return response

            with self.record_lock:
                if len(self.recorded_points) < 1:
                    response.success = False
                    response.message = 'No recorded points. Record or load trajectory first.'
                    return response

            self.stop_playback_event.clear()
            self.is_playing = True
            thread = threading.Thread(target=self._playback_worker, daemon=True)
            thread.start()

        response.success = True
        response.message = 'Playback started.'
        self.get_logger().info(response.message)
        return response

    def handle_stop_playback(self, _request, response):
        self.stop_playback_event.set()
        response.success = True
        response.message = 'Stop playback signal sent.'
        self.get_logger().info(response.message)
        return response

    def handle_clear(self, _request, response):
        self.is_recording = False
        self.stop_playback_event.set()
        with self.record_lock:
            self.recorded_points = []
            self._record_start_time = None
        response.success = True
        response.message = 'Recorded trajectory cleared.'
        self.get_logger().info(response.message)
        return response

    def handle_save(self, _request, response):
        storage_path = self.get_parameter('storage_path').value
        with self.record_lock:
            data = {
                'version': 1,
                'points': self.recorded_points,
            }

        try:
            with open(storage_path, 'w', encoding='utf-8') as f:
                yaml.safe_dump(data, f, sort_keys=False)
        except Exception as exc:
            response.success = False
            response.message = f'Save failed: {exc}'
            self.get_logger().error(response.message)
            return response

        response.success = True
        response.message = f'Trajectory saved to {storage_path}'
        self.get_logger().info(response.message)
        return response

    def handle_load(self, _request, response):
        storage_path = self.get_parameter('storage_path').value
        try:
            with open(storage_path, 'r', encoding='utf-8') as f:
                data = yaml.safe_load(f)
        except Exception as exc:
            response.success = False
            response.message = f'Load failed: {exc}'
            self.get_logger().error(response.message)
            return response

        points = data.get('points', []) if isinstance(data, dict) else []
        if not isinstance(points, list):
            response.success = False
            response.message = 'Load failed: invalid file format.'
            self.get_logger().error(response.message)
            return response

        normalized: List[Dict] = []
        for pt in points:
            try:
                normalized.append({
                    't': float(pt['t']),
                    'name': list(pt['name']),
                    'position': [float(v) for v in pt['position']],
                })
            except Exception:
                continue

        with self.record_lock:
            self.recorded_points = normalized
            self._record_start_time = None

        response.success = True
        response.message = f'Loaded {len(normalized)} points from {storage_path}'
        self.get_logger().info(response.message)
        return response

    def _playback_worker(self):
        play_speed = float(self.get_parameter('play_speed').value)
        loop_playback = bool(self.get_parameter('loop_playback').value)

        if play_speed <= 0.0:
            self.get_logger().warning('play_speed <= 0.0, force to 1.0')
            play_speed = 1.0

        try:
            while rclpy.ok():
                with self.record_lock:
                    points = list(self.recorded_points)

                if not points:
                    break

                prev_t = points[0]['t']
                for idx, pt in enumerate(points):
                    if self.stop_playback_event.is_set() or not rclpy.ok():
                        return

                    if idx > 0:
                        dt = max(0.0, (pt['t'] - prev_t) / play_speed)
                        if dt > 0.0:
                            time.sleep(dt)
                    prev_t = pt['t']

                    traj = JointTrajectory()
                    # traj.header.stamp = self.get_clock().now().to_msg()
                    traj.joint_names = pt['name']

                    point = JointTrajectoryPoint()
                    point.positions = pt['position']
                    # point.time_from_start.sec = 0
                    # point.time_from_start.nanosec = 100000000
                    traj.points.append(point)

                    self.pub.publish(traj)
                    # time.sleep(0.02)

                if not loop_playback:
                    break
        finally:
            with self.play_lock:
                self.is_playing = False
            self.stop_playback_event.clear()
            self.get_logger().info('Playback finished.')


def main(args=None):
    rclpy.init(args=args)
    node = JointTrajectoryRecorder()
    try:
        rclpy.spin(node)
    finally:
        node.stop_playback_event.set()
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
