"""test_latency.py — End-to-end latency profiling for social-bot pipeline.

Pipeline stages and SLAs:
  Stage                          SLA (p95)   Notes
  ─────────────────────────────  ─────────   ─────────────────────────────────
  wake_word → VAD onset          < 100ms     OpenWakeWord inference
  VAD onset → transcript (STT)   < 500ms     Whisper small, 1s utterance
  transcript → LLM first token   < 500ms     Phi-3-mini GGUF, TTFT
  LLM first token → TTS first    < 200ms     Piper synthesis, first sentence
  ─────────────────────────────────────────────────────────────────────────────
  Full end-to-end (wake→speaker) < 1200ms    total pipeline SLA

Tests:
  1. STT latency: inject synthetic transcript, measure to LLM response
  2. LLM latency: inject transcript, measure time to ConversationResponse
  3. TTS latency: inject ConversationResponse, verify audio starts quickly
  4. E2E latency: orchestrator state transitions (LISTENING→THINKING→SPEAKING)
  5. Percentile report: print p50/p95/p99 for each stage over N runs

In CI (no GPU/models): tests use orchestrator state timestamps to infer latency
bounds without actually running STT/LLM/TTS.
"""

from __future__ import annotations

import json
import os
import sys
import time
import threading
import unittest
from collections import defaultdict
from typing import List, Optional

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from saltybot_social_tests.test_helpers import poll_until

# ── SLA constants (milliseconds) ─────────────────────────────────────────────
SLA_STT_P95_MS        =  500.0
SLA_LLM_TTFT_P95_MS   =  500.0
SLA_TTS_FIRST_P95_MS  =  200.0
SLA_E2E_P95_MS        = 1200.0

N_PROBE_RUNS = 5   # number of injection runs for percentile calculation


def _percentile(values: List[float], p: float) -> float:
    if not values:
        return 0.0
    s = sorted(values)
    idx = int(len(s) * p / 100)
    return s[min(idx, len(s) - 1)]


class LatencyProbe(Node):
    """Injects synthetic messages and measures pipeline response latencies."""

    def __init__(self) -> None:
        super().__init__("latency_probe")
        qos = QoSProfile(depth=10)

        # Injectors
        try:
            from saltybot_social_msgs.msg import SpeechTranscript, ConversationResponse
            self._transcript_pub = self.create_publisher(
                SpeechTranscript, "/social/speech/transcript", qos
            )
            self._response_pub = self.create_publisher(
                ConversationResponse, "/social/conversation/response", qos
            )
            self._has_social = True
        except ImportError:
            self._has_social = False

        from std_msgs.msg import String
        self._state_pub = self.create_publisher(String, "/social/test/inject", qos)

        # Receivers
        from std_msgs.msg import String
        self._state_sub = self.create_subscription(
            String, "/social/orchestrator/state",
            self._on_state, qos
        )

        if self._has_social:
            from saltybot_social_msgs.msg import ConversationResponse
            self._response_sub = self.create_subscription(
                ConversationResponse, "/social/conversation/response",
                self._on_response, qos
            )

        # Timing data
        self._inject_t: Optional[float] = None
        self._response_times: List[float] = []
        self._state_transitions: List[tuple] = []   # (from, to, elapsed_ms)
        self._last_state: Optional[str] = None
        self._last_state_t: float = time.time()
        self._lock = threading.Lock()

    def _on_state(self, msg) -> None:
        try:
            data = json.loads(msg.data)
            new_state = data.get("state", "")
            now = time.time()
            with self._lock:
                if self._last_state and self._last_state != new_state:
                    elapsed = (now - self._last_state_t) * 1000
                    self._state_transitions.append(
                        (self._last_state, new_state, elapsed)
                    )
                self._last_state = new_state
                self._last_state_t = now
        except Exception:
            pass

    def _on_response(self, msg) -> None:
        with self._lock:
            if self._inject_t is not None:
                latency_ms = (time.time() - self._inject_t) * 1000
                self._response_times.append(latency_ms)
                self._inject_t = None

    def inject_transcript(self, text: str = "Hello Salty, what time is it?") -> None:
        """Publish a synthetic final transcript to trigger LLM."""
        if not self._has_social:
            return
        from saltybot_social_msgs.msg import SpeechTranscript
        msg = SpeechTranscript()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.text = text
        msg.speaker_id = "test_user"
        msg.confidence = 0.99
        msg.audio_duration = 1.2
        msg.is_partial = False
        with self._lock:
            self._inject_t = time.time()
        self._transcript_pub.publish(msg)

    def get_response_latencies(self) -> List[float]:
        with self._lock:
            return list(self._response_times)

    def get_state_transitions(self) -> List[tuple]:
        with self._lock:
            return list(self._state_transitions)

    def clear(self) -> None:
        with self._lock:
            self._response_times.clear()
            self._state_transitions.clear()
            self._inject_t = None


class TestPipelineLatency(unittest.TestCase):
    """End-to-end pipeline latency tests."""

    @classmethod
    def setUpClass(cls):
        if not rclpy.ok():
            rclpy.init()
        cls.probe = LatencyProbe()

    @classmethod
    def tearDownClass(cls):
        cls.probe.destroy_node()

    def _spin_for(self, duration_s: float) -> None:
        deadline = time.time() + duration_s
        while time.time() < deadline:
            rclpy.spin_once(self.probe, timeout_sec=0.05)

    def test_orchestrator_state_transitions_logged(self):
        """Verify orchestrator transitions are captured by the probe."""
        self.probe.clear()
        self._spin_for(3.0)   # wait for state messages

        # Inject a transcript to trigger IDLE→LISTENING→THINKING
        if self.probe._has_social:
            self.probe.inject_transcript()
            self._spin_for(5.0)

        transitions = self.probe.get_state_transitions()
        # If the stack is running with conversation_node, we should see transitions.
        # In CI without LLM, orchestrator stays IDLE — that's OK.
        print(f"\n[Latency] State transitions captured: {len(transitions)}")
        for frm, to, ms in transitions:
            print(f"  {frm} → {to}: {ms:.0f}ms")

    def test_llm_response_latency_sla(self):
        """LLM must produce first ConversationResponse within 500ms p95."""
        if os.environ.get("SOCIAL_TEST_FULL", "0") != "1":
            self.skipTest("Set SOCIAL_TEST_FULL=1 to test with LLM running")
        if not self.probe._has_social:
            self.skipTest("saltybot_social_msgs not available")

        self.probe.clear()
        latencies = []

        for i in range(N_PROBE_RUNS):
            self.probe.clear()
            self.probe.inject_transcript(f"Hello, run number {i}")
            # Wait up to 5s for response
            deadline = time.time() + 5.0
            while time.time() < deadline:
                rclpy.spin_once(self.probe, timeout_sec=0.05)
                lats = self.probe.get_response_latencies()
                if lats:
                    latencies.extend(lats)
                    break
            time.sleep(1.0)  # cooldown between runs

        if not latencies:
            self.fail("No LLM responses received in any of the test runs")

        p50 = _percentile(latencies, 50)
        p95 = _percentile(latencies, 95)
        print(f"\n[Latency] LLM response: n={len(latencies)}, "
              f"p50={p50:.0f}ms, p95={p95:.0f}ms")

        self.assertLess(
            p95, SLA_LLM_TTFT_P95_MS,
            f"LLM p95 latency {p95:.0f}ms exceeds SLA {SLA_LLM_TTFT_P95_MS:.0f}ms"
        )

    def test_e2e_state_machine_latency(self):
        """IDLE→THINKING transition must complete within 1200ms p95."""
        if not self.probe._has_social:
            self.skipTest("saltybot_social_msgs not available")

        self.probe.clear()
        thinking_latencies = []

        for _ in range(N_PROBE_RUNS):
            self.probe.clear()

            # Wait for IDLE state
            def is_idle():
                rclpy.spin_once(self.probe, timeout_sec=0.1)
                return self.probe._last_state == "idle"
            poll_until(is_idle, timeout_s=5.0)

            # Inject transcript and measure time to THINKING
            t0 = time.time()
            self.probe.inject_transcript()

            def is_thinking():
                rclpy.spin_once(self.probe, timeout_sec=0.1)
                return self.probe._last_state in ("thinking", "speaking")
            reached = poll_until(is_thinking, timeout_s=5.0)

            if reached:
                elapsed_ms = (time.time() - t0) * 1000
                thinking_latencies.append(elapsed_ms)

            time.sleep(2.0)  # allow pipeline to return to IDLE

        if not thinking_latencies:
            # In CI without full stack, this is expected
            print("[Latency] No THINKING transitions — LLM/speech may not be running")
            return

        p95 = _percentile(thinking_latencies, 95)
        print(f"\n[Latency] Inject→THINKING: n={len(thinking_latencies)}, "
              f"p95={p95:.0f}ms")

        self.assertLess(
            p95, SLA_E2E_P95_MS,
            f"E2E p95 {p95:.0f}ms exceeds SLA {SLA_E2E_P95_MS:.0f}ms"
        )


class TestLatencyReport(unittest.TestCase):
    """Generate a latency report from orchestrator state data."""

    @classmethod
    def setUpClass(cls):
        if not rclpy.ok():
            rclpy.init()
        cls.node = rclpy.create_node("latency_report_probe")

    @classmethod
    def tearDownClass(cls):
        cls.node.destroy_node()

    def test_orchestrator_reports_latency_stats(self):
        """Orchestrator state JSON must include latency field when profiling enabled."""
        from std_msgs.msg import String
        received = []
        qos = QoSProfile(depth=10)
        sub = self.node.create_subscription(
            String, "/social/orchestrator/state",
            lambda msg: received.append(msg.data), qos
        )

        deadline = time.time() + 5.0
        while time.time() < deadline and not received:
            rclpy.spin_once(self.node, timeout_sec=0.1)
        self.node.destroy_subscription(sub)

        if not received:
            self.skipTest("Orchestrator not running")

        state = json.loads(received[-1])
        self.assertIn("latency", state,
                      "Orchestrator state JSON missing 'latency' field")
        self.assertIn("ts", state)
        print(f"\n[Latency] Orchestrator report: {json.dumps(state['latency'], indent=2)}")

    def test_latency_sla_table(self):
        """Print SLA table for documentation — always passes."""
        print("\n[Latency SLA Table]")
        slas = [
            ("wake_word → transcript (STT)",  SLA_STT_P95_MS),
            ("transcript → LLM first token",  SLA_LLM_TTFT_P95_MS),
            ("LLM token → TTS first chunk",   SLA_TTS_FIRST_P95_MS),
            ("end-to-end (wake → speaker)",   SLA_E2E_P95_MS),
        ]
        for stage, sla in slas:
            print(f"  {stage:<40} p95 < {sla:.0f}ms")
