"""orchestrator_node.py — End-to-end social-bot pipeline orchestrator.

Issue #89: Main loop, GPU memory scheduling, latency profiling, watchdog.

State machine:
  IDLE → LISTENING (wake word) → THINKING (LLM generating) → SPEAKING (TTS)
       ↑___________________________________________|

Responsibilities:
  - Owns the top-level pipeline state machine
  - Monitors GPU memory and throttles inference when approaching limit
  - Profiles end-to-end latency per interaction turn
  - Health watchdog: restart speech pipeline on hang
  - Exposes /social/orchestrator/state for UI / other nodes

ROS2 topics:
  Subscribe: /social/speech/vad_state   (VadState)
             /social/speech/transcript  (SpeechTranscript)
             /social/conversation/response (ConversationResponse)
  Publish:   /social/orchestrator/state  (std_msgs/String) — JSON status blob

Parameters:
  gpu_mem_warn_mb   (float, 4000.0) — warn at this free GPU memory (MB)
  gpu_mem_throttle_mb (float, 2000.0) — suspend new inference requests
  watchdog_timeout_s  (float, 30.0) — restart speech node if no VAD in N seconds
  latency_window      (int, 20)      — rolling window for latency stats
  profile_enabled     (bool, true)
"""

from __future__ import annotations

import json
import threading
import time
from collections import deque
from enum import Enum
from typing import Optional

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

from saltybot_social_msgs.msg import VadState, SpeechTranscript, ConversationResponse


# ── State machine ──────────────────────────────────────────────────────────────

class PipelineState(Enum):
    IDLE = "idle"
    LISTENING = "listening"
    THINKING = "thinking"
    SPEAKING = "speaking"
    THROTTLED = "throttled"    # GPU OOM risk — waiting for memory


# ── Latency tracker ────────────────────────────────────────────────────────────

class LatencyTracker:
    """Rolling window latency statistics per pipeline stage."""

    def __init__(self, window: int = 20) -> None:
        self._window = window
        self._wakeword_to_transcript: deque = deque(maxlen=window)
        self._transcript_to_llm_first: deque = deque(maxlen=window)
        self._llm_first_to_tts_first: deque = deque(maxlen=window)
        self._end_to_end: deque = deque(maxlen=window)

    def record(self, stage: str, latency_ms: float) -> None:
        mapping = {
            "wakeword_to_transcript": self._wakeword_to_transcript,
            "transcript_to_llm": self._transcript_to_llm_first,
            "llm_to_tts": self._llm_first_to_tts_first,
            "end_to_end": self._end_to_end,
        }
        q = mapping.get(stage)
        if q is not None:
            q.append(latency_ms)

    def stats(self) -> dict:
        def _stat(q: deque) -> dict:
            if not q:
                return {}
            s = sorted(q)
            return {
                "mean_ms": round(sum(s) / len(s), 1),
                "p50_ms": round(s[len(s) // 2], 1),
                "p95_ms": round(s[int(len(s) * 0.95)], 1),
                "n": len(s),
            }
        return {
            "wakeword_to_transcript": _stat(self._wakeword_to_transcript),
            "transcript_to_llm": _stat(self._transcript_to_llm_first),
            "llm_to_tts": _stat(self._llm_first_to_tts_first),
            "end_to_end": _stat(self._end_to_end),
        }


# ── GPU memory helpers ────────────────────────────────────────────────────────

def get_gpu_free_mb() -> Optional[float]:
    """Return free GPU memory in MB, or None if unavailable."""
    try:
        import pycuda.driver as drv
        drv.init()
        ctx = drv.Device(0).make_context()
        free, _ = drv.mem_get_info()
        ctx.pop()
        return free / 1024 / 1024
    except Exception:
        pass
    try:
        import subprocess
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
            timeout=2
        ).decode().strip()
        return float(out.split("\n")[0])
    except Exception:
        return None


# ── Orchestrator ───────────────────────────────────────────────────────────────

class OrchestratorNode(Node):
    """Pipeline orchestrator — state machine, GPU watchdog, latency profiler."""

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

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter("gpu_mem_warn_mb", 4000.0)
        self.declare_parameter("gpu_mem_throttle_mb", 2000.0)
        self.declare_parameter("watchdog_timeout_s", 30.0)
        self.declare_parameter("latency_window", 20)
        self.declare_parameter("profile_enabled", True)
        self.declare_parameter("state_publish_rate", 2.0)   # Hz

        self._gpu_warn = self.get_parameter("gpu_mem_warn_mb").value
        self._gpu_throttle = self.get_parameter("gpu_mem_throttle_mb").value
        self._watchdog_timeout = self.get_parameter("watchdog_timeout_s").value
        self._profile = self.get_parameter("profile_enabled").value
        self._state_rate = self.get_parameter("state_publish_rate").value

        # ── Publishers / Subscribers ─────────────────────────────────────────
        qos = QoSProfile(depth=10)
        self._state_pub = self.create_publisher(String, "/social/orchestrator/state", qos)

        self._vad_sub = self.create_subscription(
            VadState, "/social/speech/vad_state", self._on_vad, qos
        )
        self._transcript_sub = self.create_subscription(
            SpeechTranscript, "/social/speech/transcript", self._on_transcript, qos
        )
        self._response_sub = self.create_subscription(
            ConversationResponse, "/social/conversation/response", self._on_response, qos
        )

        # ── State ────────────────────────────────────────────────────────────
        self._state = PipelineState.IDLE
        self._state_lock = threading.Lock()
        self._tracker = LatencyTracker(self.get_parameter("latency_window").value)

        # Timestamps for latency tracking
        self._t_wake: float = 0.0
        self._t_transcript: float = 0.0
        self._t_llm_first: float = 0.0
        self._t_tts_first: float = 0.0

        self._last_vad_t: float = time.time()
        self._last_transcript_t: float = 0.0
        self._last_response_t: float = 0.0

        # ── Timers ────────────────────────────────────────────────────────────
        self._state_timer = self.create_timer(1.0 / self._state_rate, self._publish_state)
        self._gpu_timer = self.create_timer(5.0, self._check_gpu)
        self._watchdog_timer = self.create_timer(10.0, self._watchdog)

        self.get_logger().info("OrchestratorNode ready")

    # ── State transitions ─────────────────────────────────────────────────────

    def _set_state(self, new_state: PipelineState) -> None:
        with self._state_lock:
            if self._state != new_state:
                self.get_logger().info(
                    f"Pipeline: {self._state.value} → {new_state.value}"
                )
                self._state = new_state

    # ── Subscribers ───────────────────────────────────────────────────────────

    def _on_vad(self, msg: VadState) -> None:
        self._last_vad_t = time.time()
        if msg.state == "wake_word":
            self._t_wake = time.time()
            self._set_state(PipelineState.LISTENING)
        elif msg.speech_active and self._state == PipelineState.IDLE:
            self._set_state(PipelineState.LISTENING)

    def _on_transcript(self, msg: SpeechTranscript) -> None:
        if msg.is_partial:
            return
        self._last_transcript_t = time.time()
        self._t_transcript = time.time()
        self._set_state(PipelineState.THINKING)

        if self._profile and self._t_wake > 0:
            latency_ms = (self._t_transcript - self._t_wake) * 1000
            self._tracker.record("wakeword_to_transcript", latency_ms)
            if latency_ms > 500:
                self.get_logger().warn(
                    f"Wake→transcript latency high: {latency_ms:.0f}ms (target <500ms)"
                )

    def _on_response(self, msg: ConversationResponse) -> None:
        self._last_response_t = time.time()

        if msg.is_partial and self._t_llm_first == 0.0:
            # First token from LLM
            self._t_llm_first = time.time()
            if self._profile and self._t_transcript > 0:
                latency_ms = (self._t_llm_first - self._t_transcript) * 1000
                self._tracker.record("transcript_to_llm", latency_ms)

            self._set_state(PipelineState.SPEAKING)

        if not msg.is_partial:
            # Full response complete
            self._t_llm_first = 0.0
            self._t_wake = 0.0
            self._t_transcript = 0.0

            if self._profile and self._t_wake > 0:
                e2e_ms = (time.time() - self._t_wake) * 1000
                self._tracker.record("end_to_end", e2e_ms)

            # Return to IDLE after a short delay (TTS still playing)
            threading.Timer(2.0, lambda: self._set_state(PipelineState.IDLE)).start()

    # ── GPU memory check ──────────────────────────────────────────────────────

    def _check_gpu(self) -> None:
        free_mb = get_gpu_free_mb()
        if free_mb is None:
            return

        if free_mb < self._gpu_throttle:
            self.get_logger().error(
                f"GPU memory critical: {free_mb:.0f}MB free < {self._gpu_throttle:.0f}MB threshold"
            )
            self._set_state(PipelineState.THROTTLED)
        elif free_mb < self._gpu_warn:
            self.get_logger().warn(f"GPU memory low: {free_mb:.0f}MB free")
        else:
            # Recover from throttled state
            with self._state_lock:
                if self._state == PipelineState.THROTTLED:
                    self._set_state(PipelineState.IDLE)

    # ── Watchdog ──────────────────────────────────────────────────────────────

    def _watchdog(self) -> None:
        """Alert if speech pipeline has gone silent for too long."""
        since_vad = time.time() - self._last_vad_t
        if since_vad > self._watchdog_timeout:
            self.get_logger().error(
                f"Watchdog: No VAD signal for {since_vad:.0f}s. "
                "Speech pipeline may be hung. Check /social/speech/vad_state."
            )

    # ── State publisher ───────────────────────────────────────────────────────

    def _publish_state(self) -> None:
        with self._state_lock:
            current = self._state.value

        payload = {
            "state": current,
            "ts": time.time(),
            "latency": self._tracker.stats() if self._profile else {},
            "gpu_free_mb": get_gpu_free_mb(),
        }
        msg = String()
        msg.data = json.dumps(payload)
        self._state_pub.publish(msg)


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