"""tts_node.py — Streaming TTS with Piper / first-chunk streaming.

Issue #85: Streaming TTS — Piper/XTTS integration with first-chunk streaming.

Pipeline:
  /social/conversation/response (ConversationResponse)
    → sentence split → Piper ONNX synthesis (sentence by sentence)
      → PCM16 chunks → USB speaker (sounddevice) + publish /social/tts/audio

First-chunk strategy:
  - On partial=true ConversationResponse, extract first sentence and synthesize
    immediately → audio starts before LLM finishes generating
  - On final=false, synthesize remaining sentences

Latency target: <200ms to first audio chunk.

ROS2 topics:
  Subscribe: /social/conversation/response (saltybot_social_msgs/ConversationResponse)
  Publish:   /social/tts/audio (audio_msgs/Audio or std_msgs/UInt8MultiArray fallback)

Parameters:
  voice_path      (str, "/models/piper/en_US-lessac-medium.onnx")
  sample_rate     (int, 22050)
  volume          (float, 1.0)
  audio_device    (str, "")  — sounddevice device name; "" = system default
  playback_enabled (bool, true)
  publish_audio   (bool, false)  — publish PCM to ROS2 topic
  sentence_streaming (bool, true) — synthesize sentence-by-sentence
"""

from __future__ import annotations

import queue
import threading
import time
from typing import Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile
from std_msgs.msg import UInt8MultiArray

from saltybot_social_msgs.msg import ConversationResponse
from .tts_utils import split_sentences, strip_ssml, apply_volume, chunk_pcm, estimate_duration_ms


class TtsNode(Node):
    """Streaming TTS node using Piper ONNX."""

    def __init__(self) -> None:
        super().__init__("tts_node")

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter("voice_path", "/models/piper/en_US-lessac-medium.onnx")
        self.declare_parameter("sample_rate", 22050)
        self.declare_parameter("volume", 1.0)
        self.declare_parameter("audio_device", "")
        self.declare_parameter("playback_enabled", True)
        self.declare_parameter("publish_audio", False)
        self.declare_parameter("sentence_streaming", True)

        self._voice_path = self.get_parameter("voice_path").value
        self._sample_rate = self.get_parameter("sample_rate").value
        self._volume = self.get_parameter("volume").value
        self._audio_device = self.get_parameter("audio_device").value or None
        self._playback = self.get_parameter("playback_enabled").value
        self._publish_audio = self.get_parameter("publish_audio").value
        self._sentence_streaming = self.get_parameter("sentence_streaming").value

        # ── Publishers / Subscribers ─────────────────────────────────────────
        qos = QoSProfile(depth=10)
        self._resp_sub = self.create_subscription(
            ConversationResponse, "/social/conversation/response",
            self._on_response, qos
        )
        if self._publish_audio:
            self._audio_pub = self.create_publisher(
                UInt8MultiArray, "/social/tts/audio", qos
            )

        # ── TTS engine ────────────────────────────────────────────────────────
        self._voice = None
        self._playback_queue: queue.Queue = queue.Queue(maxsize=16)
        self._current_turn = -1
        self._synthesized_turns: set = set()   # turn_ids already synthesized
        self._lock = threading.Lock()

        threading.Thread(target=self._load_voice, daemon=True).start()
        threading.Thread(target=self._playback_worker, daemon=True).start()

        self.get_logger().info(
            f"TtsNode init (voice={self._voice_path}, "
            f"streaming={self._sentence_streaming})"
        )

    # ── Voice loading ─────────────────────────────────────────────────────────

    def _load_voice(self) -> None:
        t0 = time.time()
        self.get_logger().info(f"Loading Piper voice: {self._voice_path}")
        try:
            from piper import PiperVoice
            self._voice = PiperVoice.load(self._voice_path)
            # Warmup synthesis to pre-JIT ONNX graph
            warmup_text = "Hello."
            list(self._voice.synthesize_stream_raw(warmup_text))
            self.get_logger().info(f"Piper voice ready ({time.time()-t0:.1f}s)")
        except Exception as e:
            self.get_logger().error(f"Piper voice load failed: {e}")

    # ── Response handler ──────────────────────────────────────────────────────

    def _on_response(self, msg: ConversationResponse) -> None:
        """Handle streaming LLM response — synthesize sentence by sentence."""
        if not msg.text.strip():
            return

        with self._lock:
            is_new_turn = msg.turn_id != self._current_turn
            if is_new_turn:
                self._current_turn = msg.turn_id
                # Clear old synthesized sentence cache for this new turn
                self._synthesized_turns = set()

        text = strip_ssml(msg.text)

        if self._sentence_streaming:
            sentences = split_sentences(text)
            for sentence in sentences:
                # Track which sentences we've already queued by content hash
                key = (msg.turn_id, hash(sentence))
                with self._lock:
                    if key in self._synthesized_turns:
                        continue
                    self._synthesized_turns.add(key)
                self._queue_synthesis(sentence)
        elif not msg.is_partial:
            # Non-streaming: synthesize full response at end
            self._queue_synthesis(text)

    def _queue_synthesis(self, text: str) -> None:
        """Queue a text segment for synthesis in the playback worker."""
        if not text.strip():
            return
        try:
            self._playback_queue.put_nowait(text.strip())
        except queue.Full:
            self.get_logger().warn("TTS playback queue full, dropping segment")

    # ── Playback worker ───────────────────────────────────────────────────────

    def _playback_worker(self) -> None:
        """Consume synthesis queue: synthesize → play → publish."""
        while rclpy.ok():
            try:
                text = self._playback_queue.get(timeout=0.5)
            except queue.Empty:
                continue

            if self._voice is None:
                self.get_logger().warn("TTS voice not loaded yet")
                self._playback_queue.task_done()
                continue

            t0 = time.perf_counter()
            pcm_data = self._synthesize(text)
            if pcm_data is None:
                self._playback_queue.task_done()
                continue

            synth_ms = (time.perf_counter() - t0) * 1000
            dur_ms = estimate_duration_ms(pcm_data, self._sample_rate)
            self.get_logger().debug(
                f"TTS '{text[:40]}' synth={synth_ms:.0f}ms, dur={dur_ms:.0f}ms"
            )

            if self._volume != 1.0:
                pcm_data = apply_volume(pcm_data, self._volume)

            if self._playback:
                self._play_audio(pcm_data)

            if self._publish_audio:
                self._publish_pcm(pcm_data)

            self._playback_queue.task_done()

    def _synthesize(self, text: str) -> Optional[bytes]:
        """Synthesize text to PCM16 bytes using Piper streaming."""
        if self._voice is None:
            return None
        try:
            chunks = list(self._voice.synthesize_stream_raw(text))
            return b"".join(chunks)
        except Exception as e:
            self.get_logger().error(f"TTS synthesis error: {e}")
            return None

    def _play_audio(self, pcm_data: bytes) -> None:
        """Play PCM16 data on USB speaker via sounddevice."""
        try:
            import sounddevice as sd
            import numpy as np
            samples = np.frombuffer(pcm_data, dtype=np.int16).astype(np.float32) / 32768.0
            sd.play(samples, samplerate=self._sample_rate, device=self._audio_device,
                    blocking=True)
        except Exception as e:
            self.get_logger().error(f"Audio playback error: {e}")

    def _publish_pcm(self, pcm_data: bytes) -> None:
        """Publish PCM data as UInt8MultiArray."""
        if not hasattr(self, "_audio_pub"):
            return
        msg = UInt8MultiArray()
        msg.data = list(pcm_data)
        self._audio_pub.publish(msg)


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