"""tts_node.py -- Streaming TTS with Piper / first-chunk streaming.
Issue #85/#167
"""
from __future__ import annotations
import json, queue, threading, time
from typing import Any, Dict, 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 with per-language voice switching."""
    def __init__(self) -> None:
        super().__init__("tts_node")
        self.declare_parameter("voice_path", "/models/piper/en_US-lessac-medium.onnx")
        self.declare_parameter("voice_map_json", "{}")
        self.declare_parameter("default_language", "en")
        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._voice_map_json = self.get_parameter("voice_map_json").value
        self._default_language = self.get_parameter("default_language").value or "en"
        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
        try:
            extra: Dict[str, str] = json.loads(self._voice_map_json) if self._voice_map_json.strip() not in ("{}","") else {}
        except Exception as e:
            self.get_logger().warn(f"voice_map_json parse error: {e}"); extra = {}
        self._voice_paths: Dict[str, str] = {self._default_language: self._voice_path}
        self._voice_paths.update(extra)
        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)
        self._voices: Dict[str, Any] = {}
        self._voices_lock = threading.Lock()
        self._playback_queue: queue.Queue = queue.Queue(maxsize=16)
        self._current_turn = -1
        self._synthesized_turns: set = set()
        self._lock = threading.Lock()
        threading.Thread(target=self._load_voice_for_lang, args=(self._default_language,), daemon=True).start()
        threading.Thread(target=self._playback_worker, daemon=True).start()
        self.get_logger().info(f"TtsNode init (langs={list(self._voice_paths.keys())})")

    def _load_voice_for_lang(self, lang: str) -> None:
        path = self._voice_paths.get(lang)
        if not path:
            self.get_logger().warn(f"No voice for '{lang}', fallback to '{self._default_language}'"); return
        with self._voices_lock:
            if lang in self._voices: return
        try:
            from piper import PiperVoice
            voice = PiperVoice.load(path)
            list(voice.synthesize_stream_raw("Hello."))
            with self._voices_lock: self._voices[lang] = voice
            self.get_logger().info(f"Piper [{lang}] ready")
        except Exception as e:
            self.get_logger().error(f"Piper voice load failed [{lang}]: {e}")

    def _get_voice(self, lang: str):
        with self._voices_lock:
            v = self._voices.get(lang)
            if v is not None: return v
            if lang in self._voice_paths:
                threading.Thread(target=self._load_voice_for_lang, args=(lang,), daemon=True).start()
            return self._voices.get(self._default_language)

    def _on_response(self, msg: ConversationResponse) -> None:
        if not msg.text.strip(): return
        lang = msg.language if msg.language else self._default_language
        with self._lock:
            if msg.turn_id != self._current_turn:
                self._current_turn = msg.turn_id; self._synthesized_turns = set()
        text = strip_ssml(msg.text)
        if self._sentence_streaming:
            for sentence in split_sentences(text):
                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, lang)
        elif not msg.is_partial:
            self._queue_synthesis(text, lang)

    def _queue_synthesis(self, text: str, lang: str) -> None:
        if not text.strip(): return
        try: self._playback_queue.put_nowait((text.strip(), lang))
        except queue.Full: self.get_logger().warn("TTS queue full")

    def _playback_worker(self) -> None:
        while rclpy.ok():
            try: item = self._playback_queue.get(timeout=0.5)
            except queue.Empty: continue
            text, lang = item
            voice = self._get_voice(lang)
            if voice is None:
                self.get_logger().warn(f"No voice for '{lang}'"); self._playback_queue.task_done(); continue
            t0 = time.perf_counter()
            pcm_data = self._synthesize(text, voice)
            if pcm_data is None: self._playback_queue.task_done(); continue
            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, voice) -> Optional[bytes]:
        try: return b"".join(voice.synthesize_stream_raw(text))
        except Exception as e: self.get_logger().error(f"TTS error: {e}"); return None

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

    def _publish_pcm(self, pcm_data: bytes) -> None:
        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()
