"""speech_pipeline_node.py — Wake word + VAD + Whisper STT + speaker diarization.

Issue #81: Speech pipeline for social-bot.

Pipeline:
  USB mic (PyAudio) → VAD (Silero / energy fallback)
    → wake word gate (OpenWakeWord "hey_salty")
      → utterance buffer → faster-whisper STT (Orin GPU)
        → ECAPA-TDNN speaker embedding → identify_speaker()
          → publish /social/speech/transcript + /social/speech/vad_state

Latency target: <500ms wake-to-first-token (streaming partial transcripts).

ROS2 topics published:
  /social/speech/transcript  (saltybot_social_msgs/SpeechTranscript)
  /social/speech/vad_state   (saltybot_social_msgs/VadState)

Parameters:
  mic_device_index  (int, -1 = system default)
  sample_rate       (int, 16000)
  wake_word_model   (str, "hey_salty" — OpenWakeWord model name or path)
  wake_word_threshold (float, 0.5)
  vad_threshold_db  (float, -35.0) — fallback energy VAD
  use_silero_vad    (bool, true)
  whisper_model     (str, "small")
  whisper_compute_type (str, "float16")
  speaker_threshold (float, 0.65) — cosine similarity threshold
  speaker_db_path   (str, "/social_db/speaker_embeddings.json")
  publish_partial   (bool, true)
"""

from __future__ import annotations

import json
import os
import threading
import time
from typing import Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, QoSReliabilityPolicy, QoSDurabilityPolicy
from builtin_interfaces.msg import Time
from std_msgs.msg import Header

from saltybot_social_msgs.msg import SpeechTranscript, VadState
from .speech_utils import (
    EnergyVad, UtteranceSegmenter,
    pcm16_to_float32, rms_db, identify_speaker,
    SAMPLE_RATE, CHUNK_SAMPLES,
)


class SpeechPipelineNode(Node):
    """Wake word → VAD → STT → diarization → ROS2 publisher."""

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

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter("mic_device_index", -1)
        self.declare_parameter("sample_rate", SAMPLE_RATE)
        self.declare_parameter("wake_word_model", "hey_salty")
        self.declare_parameter("wake_word_threshold", 0.5)
        self.declare_parameter("vad_threshold_db", -35.0)
        self.declare_parameter("use_silero_vad", True)
        self.declare_parameter("whisper_model", "small")
        self.declare_parameter("whisper_compute_type", "float16")
        self.declare_parameter("whisper_language", "")
        self.declare_parameter("speaker_threshold", 0.65)
        self.declare_parameter("speaker_db_path", "/social_db/speaker_embeddings.json")
        self.declare_parameter("publish_partial", True)

        self._mic_idx = self.get_parameter("mic_device_index").value
        self._rate = self.get_parameter("sample_rate").value
        self._ww_model = self.get_parameter("wake_word_model").value
        self._ww_thresh = self.get_parameter("wake_word_threshold").value
        self._vad_thresh_db = self.get_parameter("vad_threshold_db").value
        self._use_silero = self.get_parameter("use_silero_vad").value
        self._whisper_model_name = self.get_parameter("whisper_model").value
        self._compute_type = self.get_parameter("whisper_compute_type").value
        self._whisper_language = self.get_parameter("whisper_language").value or None
        self._speaker_thresh = self.get_parameter("speaker_threshold").value
        self._speaker_db = self.get_parameter("speaker_db_path").value
        self._publish_partial = self.get_parameter("publish_partial").value

        # ── Publishers ───────────────────────────────────────────────────────
        qos = QoSProfile(depth=10)
        self._transcript_pub = self.create_publisher(SpeechTranscript,
                                                     "/social/speech/transcript", qos)
        self._vad_pub = self.create_publisher(VadState, "/social/speech/vad_state", qos)

        # ── Internal state ───────────────────────────────────────────────────
        self._whisper = None
        self._ecapa = None
        self._oww = None
        self._vad = None
        self._segmenter = None
        self._known_speakers: dict = {}
        self._audio_stream = None
        self._pa = None
        self._lock = threading.Lock()
        self._running = False
        self._wake_word_active = False
        self._wake_expiry = 0.0
        self._wake_window_s = 8.0   # seconds to listen after wake word
        self._turn_counter = 0

        # ── Lazy model loading in background thread ──────────────────────────
        self._model_thread = threading.Thread(target=self._load_models, daemon=True)
        self._model_thread.start()

    # ── Model loading ─────────────────────────────────────────────────────────

    def _load_models(self) -> None:
        """Load all AI models (runs in background to not block ROS spin)."""
        self.get_logger().info("Loading speech models (Whisper, ECAPA-TDNN, VAD)...")
        t0 = time.time()

        # faster-whisper
        try:
            from faster_whisper import WhisperModel
            self._whisper = WhisperModel(
                self._whisper_model_name,
                device="cuda",
                compute_type=self._compute_type,
                download_root="/models",
            )
            self.get_logger().info(
                f"Whisper '{self._whisper_model_name}' loaded "
                f"({time.time()-t0:.1f}s)"
            )
        except Exception as e:
            self.get_logger().error(f"Whisper load failed: {e}")

        # Silero VAD
        if self._use_silero:
            try:
                from silero_vad import load_silero_vad
                self._vad = load_silero_vad()
                self.get_logger().info("Silero VAD loaded")
            except Exception as e:
                self.get_logger().warn(f"Silero VAD unavailable ({e}), using energy VAD")
                self._vad = None

        energy_vad = EnergyVad(threshold_db=self._vad_thresh_db)
        self._segmenter = UtteranceSegmenter(energy_vad)

        # ECAPA-TDNN speaker embeddings
        try:
            from speechbrain.pretrained import EncoderClassifier
            self._ecapa = EncoderClassifier.from_hparams(
                source="speechbrain/spkrec-ecapa-voxceleb",
                savedir="/models/speechbrain_ecapa",
            )
            self.get_logger().info("ECAPA-TDNN loaded")
        except Exception as e:
            self.get_logger().warn(f"ECAPA-TDNN unavailable ({e}), using 'unknown' speaker")

        # OpenWakeWord
        try:
            import openwakeword
            from openwakeword.model import Model as OWWModel
            self._oww = OWWModel(
                wakeword_models=[self._ww_model],
                inference_framework="onnx",
            )
            self.get_logger().info(f"OpenWakeWord '{self._ww_model}' loaded")
        except Exception as e:
            self.get_logger().warn(
                f"OpenWakeWord unavailable ({e}). Robot will listen continuously."
            )

        # Load speaker DB
        self._load_speaker_db()

        # Start audio capture
        self._start_audio()
        self.get_logger().info(f"Speech pipeline ready ({time.time()-t0:.1f}s total)")

    def _load_speaker_db(self) -> None:
        """Load speaker embedding database from JSON."""
        if os.path.exists(self._speaker_db):
            try:
                with open(self._speaker_db) as f:
                    self._known_speakers = json.load(f)
                self.get_logger().info(
                    f"Speaker DB loaded: {len(self._known_speakers)} persons"
                )
            except Exception as e:
                self.get_logger().error(f"Speaker DB load error: {e}")

    # ── Audio capture ─────────────────────────────────────────────────────────

    def _start_audio(self) -> None:
        """Open PyAudio stream and start capture thread."""
        try:
            import pyaudio
            self._pa = pyaudio.PyAudio()
            dev_idx = None if self._mic_idx < 0 else self._mic_idx
            self._audio_stream = self._pa.open(
                format=pyaudio.paInt16,
                channels=1,
                rate=self._rate,
                input=True,
                input_device_index=dev_idx,
                frames_per_buffer=CHUNK_SAMPLES,
            )
            self._running = True
            t = threading.Thread(target=self._audio_loop, daemon=True)
            t.start()
            self.get_logger().info(
                f"Audio capture started (device={self._mic_idx}, {self._rate}Hz)"
            )
        except Exception as e:
            self.get_logger().error(f"Audio stream error: {e}")

    def _audio_loop(self) -> None:
        """Continuous audio capture loop — runs in dedicated thread."""
        while self._running:
            try:
                raw = self._audio_stream.read(CHUNK_SAMPLES, exception_on_overflow=False)
            except Exception:
                continue

            samples = pcm16_to_float32(raw)
            db = rms_db(samples)
            now = time.time()

            # VAD state publish (10Hz — every 3 chunks at 30ms each)
            vad_active = self._check_vad(samples)
            state_str = "silence"

            # Wake word check
            if self._oww is not None:
                preds = self._oww.predict(samples)
                score = preds.get(self._ww_model, 0.0)
                if isinstance(score, (list, tuple)):
                    score = score[-1]
                if score >= self._ww_thresh:
                    self._wake_word_active = True
                    self._wake_expiry = now + self._wake_window_s
                    state_str = "wake_word"
                    self.get_logger().info(
                        f"Wake word detected (score={score:.2f})"
                    )
            else:
                # No wake word detector — always listen
                self._wake_word_active = True
                self._wake_expiry = now + self._wake_window_s

            if self._wake_word_active and now > self._wake_expiry:
                self._wake_word_active = False

            # Publish VAD state
            if vad_active:
                state_str = "speech"
            self._publish_vad(vad_active, db, state_str)

            # Only feed segmenter when wake word is active
            if not self._wake_word_active:
                continue

            if self._segmenter is None:
                continue

            completed = self._segmenter.push(samples)
            for utt_samples, duration in completed:
                threading.Thread(
                    target=self._process_utterance,
                    args=(utt_samples, duration),
                    daemon=True,
                ).start()

    def _check_vad(self, samples: list) -> bool:
        """Run Silero VAD or fall back to segmenter's energy VAD."""
        if self._vad is not None:
            try:
                import torch
                t = torch.tensor(samples, dtype=torch.float32)
                prob = self._vad(t, self._rate).item()
                return prob > 0.5
            except Exception:
                pass
        if self._segmenter is not None:
            return self._segmenter._vad.is_active
        return False

    # ── STT + diarization ─────────────────────────────────────────────────────

    def _process_utterance(self, samples: list, duration: float) -> None:
        """Transcribe utterance and identify speaker. Runs in thread."""
        if self._whisper is None:
            return

        t0 = time.perf_counter()

        # Convert to numpy for faster-whisper
        try:
            import numpy as np
            audio_np = np.array(samples, dtype=np.float32)
        except ImportError:
            audio_np = samples

        # Speaker embedding first (can be concurrent with transcription)
        speaker_id = "unknown"
        if self._ecapa is not None:
            try:
                import torch
                audio_tensor = torch.tensor([samples], dtype=torch.float32)
                with torch.no_grad():
                    emb = self._ecapa.encode_batch(audio_tensor)
                emb_list = emb[0].cpu().numpy().tolist()
                speaker_id = identify_speaker(
                    emb_list, self._known_speakers, self._speaker_thresh
                )
            except Exception as e:
                self.get_logger().debug(f"Speaker ID error: {e}")

        # Streaming Whisper transcription with language detection
        partial_text = ""
        detected_lang = self._whisper_language or "en"
        try:
            segments_gen, info = self._whisper.transcribe(
                audio_np,
                language=self._whisper_language,  # None = auto-detect
                beam_size=3,
                vad_filter=False,
            )
            if hasattr(info, "language") and info.language:
                detected_lang = info.language
            for seg in segments_gen:
                partial_text += seg.text.strip() + " "
                if self._publish_partial:
                    self._publish_transcript(
                        partial_text.strip(), speaker_id, 0.0, duration,
                        language=detected_lang, is_partial=True,
                    )
        except Exception as e:
            self.get_logger().error(f"Whisper error: {e}")
            return

        final_text = partial_text.strip()
        if not final_text:
            return

        latency_ms = (time.perf_counter() - t0) * 1000
        self.get_logger().info(
            f"STT [{speaker_id}/{detected_lang}] ({duration:.1f}s, {latency_ms:.0f}ms): "
            f"'{final_text}'"
        )
        self._publish_transcript(
            final_text, speaker_id, 0.9, duration,
            language=detected_lang, is_partial=False,
        )

    # ── Publishers ────────────────────────────────────────────────────────────

    def _publish_transcript(
        self, text: str, speaker_id: str, confidence: float,
        duration: float, language: str = "en", is_partial: bool = False,
    ) -> None:
        msg = SpeechTranscript()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.text = text
        msg.speaker_id = speaker_id
        msg.confidence = confidence
        msg.audio_duration = duration
        msg.language = language
        msg.is_partial = is_partial
        self._transcript_pub.publish(msg)

    def _publish_vad(self, speech_active: bool, db: float, state: str) -> None:
        msg = VadState()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.speech_active = speech_active
        msg.energy_db = db
        msg.state = state
        self._vad_pub.publish(msg)

    # ── Cleanup ───────────────────────────────────────────────────────────────

    def destroy_node(self) -> None:
        self._running = False
        if self._audio_stream:
            self._audio_stream.stop_stream()
            self._audio_stream.close()
        if self._pa:
            self._pa.terminate()
        super().destroy_node()


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