"""
audio_scene_node.py — Audio scene classifier node (Issue #353).

Buffers raw PCM audio from the audio_common stack, runs a nearest-centroid
MFCC classifier at 1 Hz, and publishes the estimated environment label.

Subscribes
----------
  /audio/audio       audio_common_msgs/AudioData     (BEST_EFFORT)
  /audio/audio_info  audio_common_msgs/AudioInfo     (RELIABLE, latched)

Publishes
---------
  /saltybot/audio_scene   saltybot_scene_msgs/AudioScene   (1 Hz)

Parameters
----------
  sample_rate       int    16000   Expected audio sample rate (Hz)
  channels          int    1       Expected number of channels (mono=1)
  clip_duration_s   float  1.0     Duration of each classification window (s)
  publish_rate_hz   float  1.0     Classification + publish rate (Hz)

Notes
-----
  * Audio data is expected as signed 16-bit little-endian PCM ('S16LE').
    If AudioInfo arrives first with a different format, a warning is logged.
  * If the audio buffer does not contain enough samples when the timer fires,
    the existing buffer (padded with zeros to clip_duration_s) is used.
  * The publish rate is independent of the audio callback rate; bursts and
    gaps are handled gracefully.
"""

from __future__ import annotations

import struct
from collections import deque
from typing import Deque

import numpy as np

import rclpy
from rclpy.node import Node
from rclpy.qos import (
    QoSProfile,
    ReliabilityPolicy,
    HistoryPolicy,
    DurabilityPolicy,
)

try:
    from audio_common_msgs.msg import AudioData, AudioInfo
    _HAVE_AUDIO_MSGS = True
except ImportError:
    _HAVE_AUDIO_MSGS = False

from saltybot_scene_msgs.msg import AudioScene
from ._audio_scene import classify_audio, _N_FEATURES


_SENSOR_QOS = QoSProfile(
    reliability=ReliabilityPolicy.BEST_EFFORT,
    history=HistoryPolicy.KEEP_LAST,
    depth=10,
)
_LATCHED_QOS = QoSProfile(
    reliability=ReliabilityPolicy.RELIABLE,
    history=HistoryPolicy.KEEP_LAST,
    depth=1,
    durability=DurabilityPolicy.TRANSIENT_LOCAL,
)


class AudioSceneNode(Node):

    def __init__(self) -> None:
        super().__init__('audio_scene_node')

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter('sample_rate',     16_000)
        self.declare_parameter('channels',        1)
        self.declare_parameter('clip_duration_s', 1.0)
        self.declare_parameter('publish_rate_hz', 1.0)

        p = self.get_parameter
        self._sr           = int(p('sample_rate').value)
        self._channels     = int(p('channels').value)
        self._clip_dur     = float(p('clip_duration_s').value)
        self._pub_rate     = float(p('publish_rate_hz').value)

        self._clip_samples = int(self._sr * self._clip_dur)

        # ── Audio sample ring buffer ─────────────────────────────────────────
        # Store float32 mono samples; deque with max size = 2× clip
        self._buffer: Deque[float] = deque(
            maxlen=self._clip_samples * 2
        )
        self._audio_format = 'S16LE'   # updated from AudioInfo if available

        # ── Subscribers ──────────────────────────────────────────────────────
        if _HAVE_AUDIO_MSGS:
            self._audio_sub = self.create_subscription(
                AudioData, '/audio/audio',
                self._on_audio, _SENSOR_QOS,
            )
            self._info_sub = self.create_subscription(
                AudioInfo, '/audio/audio_info',
                self._on_audio_info, _LATCHED_QOS,
            )
        else:
            self.get_logger().warn(
                'audio_common_msgs not available — no audio will be received. '
                'Install ros-humble-audio-common-msgs to enable audio input.'
            )

        # ── Publisher + timer ────────────────────────────────────────────────
        self._pub = self.create_publisher(
            AudioScene, '/saltybot/audio_scene', 10
        )
        period = 1.0 / max(self._pub_rate, 0.01)
        self._timer = self.create_timer(period, self._on_timer)

        self.get_logger().info(
            f'audio_scene_node ready — '
            f'sr={self._sr} Hz, clip={self._clip_dur:.1f}s, '
            f'rate={self._pub_rate:.1f} Hz'
        )

    # ── Callbacks ─────────────────────────────────────────────────────────────

    def _on_audio_info(self, msg) -> None:
        """Latch audio format from AudioInfo."""
        fmt = getattr(msg, 'sample_format', 'S16LE') or 'S16LE'
        if fmt != self._audio_format:
            self.get_logger().info(f'Audio format updated: {fmt!r}')
            self._audio_format = fmt
        reported_sr = getattr(msg, 'sample_rate', self._sr)
        if reported_sr and reported_sr != self._sr:
            self.get_logger().warn(
                f'AudioInfo sample_rate={reported_sr} != '
                f'param sample_rate={self._sr} — using param value',
                throttle_duration_sec=30.0,
            )

    def _on_audio(self, msg: 'AudioData') -> None:
        """Decode raw PCM bytes and push float32 mono samples to buffer."""
        raw = bytes(msg.data)
        if len(raw) < 2:
            return

        try:
            n_samples = len(raw) // 2
            samples_i16 = struct.unpack(f'<{n_samples}h', raw[:n_samples * 2])
        except struct.error as exc:
            self.get_logger().warn(
                f'Failed to decode audio frame: {exc}',
                throttle_duration_sec=5.0,
            )
            return

        # Convert to float32 mono
        arr = np.array(samples_i16, dtype=np.float32) / 32768.0
        if self._channels > 1:
            # Downmix to mono: take channel 0
            arr = arr[::self._channels]

        self._buffer.extend(arr.tolist())

    def _on_timer(self) -> None:
        """Classify buffered audio and publish result."""
        n_buf = len(self._buffer)
        if n_buf >= self._clip_samples:
            # Take the most recent clip_samples
            samples = np.array(
                list(self._buffer)[-self._clip_samples:],
                dtype=np.float64,
            )
        elif n_buf > 0:
            # Pad with zeros at the start
            samples = np.zeros(self._clip_samples, dtype=np.float64)
            buf_arr = np.array(list(self._buffer), dtype=np.float64)
            samples[-n_buf:] = buf_arr
        else:
            # No audio received — publish silence (indoor default)
            samples = np.zeros(self._clip_samples, dtype=np.float64)

        result = classify_audio(samples, sr=self._sr)

        msg = AudioScene()
        msg.header.stamp    = self.get_clock().now().to_msg()
        msg.header.frame_id = 'audio'
        msg.label           = result.label
        msg.confidence      = float(result.confidence)
        feat_padded         = np.zeros(_N_FEATURES, dtype=np.float32)
        n = min(len(result.features), _N_FEATURES)
        feat_padded[:n]     = result.features[:n].astype(np.float32)
        msg.features        = feat_padded.tolist()

        self._pub.publish(msg)
        self.get_logger().debug(
            f'audio_scene: {result.label!r} conf={result.confidence:.2f}'
        )


def main(args=None) -> None:
    rclpy.init(args=args)
    node = AudioSceneNode()
    try:
        rclpy.spin(node)
    finally:
        node.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()
