#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSHistoryPolicy, QoSProfile, QoSReliabilityPolicy
from ros2_plugin_proto.msg import RosMsgWrapper

from aimdk.protocol_pb2 import ProcessedAudioOutput, AudioVADState

import datetime
import os

class AudioSubscriber(Node):
    def __init__(self):
        super().__init__("audio_subscriber")

        # 音频缓冲区，按 stream_id 分别存储
        self.audio_buffers = {}  # {stream_id: bytearray()}
        self.recording_state = {}  # {stream_id: bool} 记录是否正在录音

        # 创建音频文件存储目录
        self.audio_output_dir = "audio_recordings"
        os.makedirs(self.audio_output_dir, exist_ok=True)

        qos_profile = QoSProfile(
            history=QoSHistoryPolicy.KEEP_LAST,
            depth=10,
            reliability=QoSReliabilityPolicy.BEST_EFFORT,
        )

        self.subscription = self.create_subscription(
            RosMsgWrapper,
            "/agent/process_audio_output/pb_3Aaimdk_2Eprotocol_2EProcessedAudioOutput",
            self.audio_callback,
            qos_profile,
        )

        self.get_logger().info("Started subscribing to noise-reduced audio data...")

    def audio_callback(self, msg):
        try:
            # Check if serialization type is pb
            if msg.serialization_type != "pb":
                self.get_logger().warn(f"Unsupported serialization type: {msg.serialization_type}")
                return

            # Convert the data field from list[bytes] to bytes
            audio_data_bytes = b"".join(msg.data)

            # Parse the message using the generated protobuf class
            processed_audio = ProcessedAudioOutput()
            processed_audio.ParseFromString(audio_data_bytes)

            self.get_logger().info(
                f"Received audio data: stream_id={processed_audio.stream_id}, "
                f"vad_state={processed_audio.vad_state}, "
                f"audio_size={len(processed_audio.audio_data)} bytes"
            )

            # Handle audio based on VAD state
            self.handle_vad_state(processed_audio)

        except Exception as e:
            self.get_logger().error(f"Error processing audio message: {e}")

    def handle_vad_state(self, processed_audio):
        """Handle different VAD states"""
        vad_state = processed_audio.vad_state
        stream_id = processed_audio.stream_id
        audio_data = processed_audio.audio_data

        # Initialize the buffer for this stream_id if it does not exist
        if stream_id not in self.audio_buffers:
            self.audio_buffers[stream_id] = bytearray()
            self.recording_state[stream_id] = False

        # VAD state name mapping
        vad_state_names = {
            AudioVADState.AUDIO_VAD_STATE_NONE: "No Voice",
            AudioVADState.AUDIO_VAD_STATE_BEGIN: "Voice Start",
            AudioVADState.AUDIO_VAD_STATE_PROCESSING: "Voice Processing",
            AudioVADState.AUDIO_VAD_STATE_END: "Voice End",
        }

        stream_names = {1: "Built-in Microphone", 2: "External Microphone"}

        self.get_logger().info(
            f"[{stream_names.get(stream_id, f'Unknown Stream {stream_id}')}] "
            f"VAD State: {vad_state_names.get(vad_state, f'Unknown State {vad_state}')} "
            f"Audio Data: {len(audio_data)} bytes"
        )

        # Handle audio data based on VAD state
        if vad_state == AudioVADState.AUDIO_VAD_STATE_BEGIN:
            self.get_logger().info("🎤 Voice start detected")
            # Start a new recording, clear the buffer
            self.audio_buffers[stream_id].clear()
            self.recording_state[stream_id] = True
            # Add current audio data
            if len(audio_data) > 0:
                self.audio_buffers[stream_id].extend(audio_data)

        elif vad_state == AudioVADState.AUDIO_VAD_STATE_PROCESSING:
            self.get_logger().info("🔄 Voice processing...")
            # If recording, continue adding audio data to the buffer
            if self.recording_state[stream_id] and len(audio_data) > 0:
                self.audio_buffers[stream_id].extend(audio_data)

        elif vad_state == AudioVADState.AUDIO_VAD_STATE_END:
            self.get_logger().info("✅ Voice end")
            # Add the last audio data
            if self.recording_state[stream_id] and len(audio_data) > 0:
                self.audio_buffers[stream_id].extend(audio_data)

            # Save the complete audio segment
            if (
                self.recording_state[stream_id]
                and len(self.audio_buffers[stream_id]) > 0
            ):
                self.save_audio_segment(bytes(self.audio_buffers[stream_id]), stream_id)

            # End recording
            self.recording_state[stream_id] = False

        elif vad_state == AudioVADState.AUDIO_VAD_STATE_NONE:
            # No voice state, do not record
            if self.recording_state[stream_id]:
                self.get_logger().info("⏹️ Recording state reset")
                self.recording_state[stream_id] = False

        # Output current buffer state
        if stream_id in self.audio_buffers:
            buffer_size = len(self.audio_buffers[stream_id])
            recording = self.recording_state[stream_id]
            self.get_logger().debug(
                f"[Stream {stream_id}] 缓冲区大小：{buffer_size} bytes, 录音状态：{recording}"
            )

    def save_audio_segment(self, audio_data, stream_id):
        """Save audio segment 16kHz, 16-bit, mono PCM"""
        if len(audio_data) > 0:
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")

            # Create subdirectory by stream_id
            stream_dir = os.path.join(self.audio_output_dir, f"stream_{stream_id}")
            os.makedirs(stream_dir, exist_ok=True)

            # Generate file name
            stream_names = {1: "internal_mic", 2: "external_mic"}
            stream_name = stream_names.get(stream_id, f"stream_{stream_id}")
            filename = f"{stream_name}_{timestamp}.pcm"
            filepath = os.path.join(stream_dir, filename)

            try:
                with open(filepath, "wb") as f:
                    f.write(audio_data)

                self.get_logger().info(
                    f"Audio segment saved: {filepath} (size: {len(audio_data)} bytes)"
                )

                # Record the duration of the audio file (assuming 16kHz, 16-bit, mono)
                sample_rate = 16000
                bits_per_sample = 16
                channels = 1
                bytes_per_sample = bits_per_sample // 8
                total_samples = len(audio_data) // (bytes_per_sample * channels)
                duration_seconds = total_samples / sample_rate

                self.get_logger().info(
                    f"Audio duration: {duration_seconds:.2f} seconds ({total_samples} samples)"
                )

            except Exception as e:
                self.get_logger().error(f"Failed to save audio file: {e}")

    def get_buffer_info(self):
        """Get information of all buffers (for debugging)"""
        info = {}
        for stream_id in self.audio_buffers:
            info[stream_id] = {
                "buffer_size": len(self.audio_buffers[stream_id]),
                "recording": self.recording_state[stream_id],
            }
        return info

def main(args=None):
    rclpy.init(args=args)

    audio_subscriber = AudioSubscriber()

    try:
        audio_subscriber.get_logger().info("Listening to noise-reduced audio data, press Ctrl+C to exit...")
        rclpy.spin(audio_subscriber)
    except KeyboardInterrupt:
        audio_subscriber.get_logger().info("Received exit signal, shutting down...")
    finally:
        audio_subscriber.destroy_node()
        rclpy.shutdown()

if __name__ == "__main__":
    main()

export ROS_DOMAIN_ID=232
export FASTRTPS_DEFAULT_PROFILES_FILE=/agibot/software/v0/entry/cfg/ros_dds_configuration.xml

ros2 service call /audio_5Fmsgs/srv/RequestAudioFocus audio_msgs/srv/RequestAudioFocus "{focus_requester: {pkg_name: audio_examples_sender, priority: 6, priority_weight: 1}}"

#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from builtin_interfaces.msg import Time as TimeMsg
from audio_msgs.msg import AudioPlayback,AudioInfo,AudioData  # 需确保 audio msg 可导入

class AudioPlaybackPublisher(Node):
    def __init__(self):
        super().__init__('audio_playback_publisher')
        self.pub = self.create_publisher(AudioPlayback, '/audiohal/audio/playback', 10)
        self.timer = self.create_timer(1.0, self.timer_callback)  # 每秒发送

    def timer_callback(self):
        msg = AudioPlayback()

        # 填写时间戳（字段名 attachments 中为 stamps）
        now = self.get_clock().now().to_msg()
        # 兼容性：优先使用 stamps，否则尝试 stamp
        if hasattr(msg, 'stamps'):
            msg.stamps = now
        elif hasattr(msg, 'stamp'):
            msg.stamp = now

        # 填写 info（AudioInfo）
        info = AudioInfo()
        info.channels = 1          # 单声道
        info.sample_rate = 16000   # 16kHz
        info.sample_format = 'S16LE'
        info.coding_format = 'pcm'
        msg.info = info

        # 填写 data（AudioData），示例给一小段原始 bytes
        sample_bytes = b'\x00\x01\x02\x03\x04\x05\x06\x07'  # 示例数据
        msg.data = AudioData(data=list(sample_bytes))

        msg.pkg_name = 'audio_examples_sender'
        msg.token_id = 'token-0001'

        self.pub.publish(msg)
        self.get_logger().info(
            f'Published AudioPlayback pkg={msg.pkg_name} token={msg.token_id} '
            f'data_len={len(msg.data.data)} sample_rate={msg.info.sample_rate}'
        )


def main(args=None):
    rclpy.init(args=args)
    node = AudioPlaybackPublisher()
    try:
        rclpy.spin(node)
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()

#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from rclpy.qos import (
    QoSDurabilityPolicy,
    QoSHistoryPolicy,
    QoSProfile,
    QoSReliabilityPolicy,
)
from ros2_plugin_proto.msg import RosMsgWrapper

from aimdk.protocol_pb2 import WakeUpResult


TOPIC = "/agent/wakeup/pb_3Aaimdk_2Eprotocol_2EWakeUpResult"


class WakeUpSubscriber(Node):
    def __init__(self):
        super().__init__("wakeup_result_subscriber")

        qos_profile = QoSProfile(
            history=QoSHistoryPolicy.KEEP_LAST,
            depth=10,
            reliability=QoSReliabilityPolicy.RELIABLE,
            durability=QoSDurabilityPolicy.VOLATILE,
        )

        self.subscription = self.create_subscription(
            RosMsgWrapper,
            TOPIC,
            self.wakeup_callback,
            qos_profile,
        )

        self.get_logger().info(f"Started subscribing to WakeUpResult: {TOPIC}")

    def wakeup_callback(self, msg):
        try:
            if msg.serialization_type != "pb":
                self.get_logger().warn(
                    f"Unsupported serialization type: {msg.serialization_type}"
                )
                return

            # Concatenate bytes
            raw_bytes = b"".join(msg.data)

            # Parse protobuf
            wakeup_result = WakeUpResult()
            wakeup_result.ParseFromString(raw_bytes)

            # Log output
            import json
            from google.protobuf.json_format import MessageToDict

            self.get_logger().info(
                f"WakeUpResult: {json.dumps(MessageToDict(wakeup_result, preserving_proto_field_name=True), ensure_ascii=False, indent=2)}"
            )

        except Exception as e:
            self.get_logger().error(f"Error parsing WakeUpResult data: {e}")


def main(args=None):
    rclpy.init(args=args)
    node = WakeUpSubscriber()

    try:
        node.get_logger().info("Listening to WakeUpResult, press Ctrl+C to exit...")
        rclpy.spin(node)
    except KeyboardInterrupt:
        node.get_logger().info("Exiting...")
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == "__main__":
    main()

export ROS_DOMAIN_ID=232
export FASTRTPS_DEFAULT_PROFILES_FILE=/agibot/software/v0/entry/cfg/ros_dds_configuration.xml