"""conversation_node.py — Local LLM conversation engine with per-person context.

Issue #83: Conversation engine for social-bot.

Stack: Phi-3-mini or Llama-3.2-3B GGUF Q4_K_M via llama-cpp-python (CUDA).
Subscribes /social/speech/transcript → builds per-person prompt → streams
token output → publishes /social/conversation/response.

Streaming: publishes partial=true tokens as they arrive, then final=false
at end of generation. TTS node can begin synthesis on first sentence boundary.

ROS2 topics:
  Subscribe:  /social/speech/transcript (saltybot_social_msgs/SpeechTranscript)
  Publish:    /social/conversation/response (saltybot_social_msgs/ConversationResponse)

Parameters:
  model_path       (str, "/models/gguf/phi-3-mini-4k-instruct-q4_k_m.gguf")
  n_ctx            (int, 4096)
  n_gpu_layers     (int, 20)    — GPU offload layers (increase for more VRAM usage)
  max_tokens       (int, 200)
  temperature      (float, 0.7)
  top_p            (float, 0.9)
  soul_path        (str, "/soul/SOUL.md")
  context_db_path  (str, "/social_db/conversation_context.json")
  save_interval_s  (float, 30.0) — how often to persist context to disk
  stream           (bool, true)
"""

from __future__ import annotations

import threading
import time
from typing import Optional

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

from saltybot_social_msgs.msg import SpeechTranscript, ConversationResponse
from .llm_context import ContextStore, build_llama_prompt, load_system_prompt, needs_summary_prompt


class ConversationNode(Node):
    """Local LLM inference node with per-person conversation memory."""

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

        # ── Parameters ──────────────────────────────────────────────────────
        self.declare_parameter("model_path",
                               "/models/gguf/phi-3-mini-4k-instruct-q4_k_m.gguf")
        self.declare_parameter("n_ctx", 4096)
        self.declare_parameter("n_gpu_layers", 20)
        self.declare_parameter("max_tokens", 200)
        self.declare_parameter("temperature", 0.7)
        self.declare_parameter("top_p", 0.9)
        self.declare_parameter("soul_path", "/soul/SOUL.md")
        self.declare_parameter("context_db_path", "/social_db/conversation_context.json")
        self.declare_parameter("save_interval_s", 30.0)
        self.declare_parameter("stream", True)

        self._model_path = self.get_parameter("model_path").value
        self._n_ctx = self.get_parameter("n_ctx").value
        self._n_gpu = self.get_parameter("n_gpu_layers").value
        self._max_tokens = self.get_parameter("max_tokens").value
        self._temperature = self.get_parameter("temperature").value
        self._top_p = self.get_parameter("top_p").value
        self._soul_path = self.get_parameter("soul_path").value
        self._db_path = self.get_parameter("context_db_path").value
        self._save_interval = self.get_parameter("save_interval_s").value
        self._stream = self.get_parameter("stream").value

        # ── Publishers / Subscribers ─────────────────────────────────────────
        qos = QoSProfile(depth=10)
        self._resp_pub = self.create_publisher(
            ConversationResponse, "/social/conversation/response", qos
        )
        self._transcript_sub = self.create_subscription(
            SpeechTranscript, "/social/speech/transcript",
            self._on_transcript, qos
        )

        # ── State ────────────────────────────────────────────────────────────
        self._llm = None
        self._system_prompt = load_system_prompt(self._soul_path)
        self._ctx_store = ContextStore(self._db_path)
        self._lock = threading.Lock()
        self._turn_counter = 0
        self._generating = False
        self._last_save = time.time()

        # ── Load LLM in background ────────────────────────────────────────────
        threading.Thread(target=self._load_llm, daemon=True).start()

        # ── Periodic context save ────────────────────────────────────────────
        self._save_timer = self.create_timer(self._save_interval, self._save_context)

        # ── Emotion context (Issue #161) ──────────────────────────────────────
        self._emotions: Dict[int, str] = {}
        self.create_subscription(
            String, "/social/emotion/context",
            self._on_emotion_context, 10
        )

        self.get_logger().info(
            f"ConversationNode init (model={self._model_path}, "
            f"gpu_layers={self._n_gpu}, ctx={self._n_ctx})"
        )

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

    def _load_llm(self) -> None:
        t0 = time.time()
        self.get_logger().info(f"Loading LLM: {self._model_path}")
        try:
            from llama_cpp import Llama
            self._llm = Llama(
                model_path=self._model_path,
                n_ctx=self._n_ctx,
                n_gpu_layers=self._n_gpu,
                n_threads=4,
                verbose=False,
            )
            self.get_logger().info(
                f"LLM ready ({time.time()-t0:.1f}s). "
                f"Context: {self._n_ctx} tokens, GPU layers: {self._n_gpu}"
            )
        except Exception as e:
            self.get_logger().error(f"LLM load failed: {e}")

    # ── Transcript callback ───────────────────────────────────────────────────

    def _on_transcript(self, msg: SpeechTranscript) -> None:
        """Handle final transcripts only (skip streaming partials)."""
        if msg.is_partial:
            return
        if not msg.text.strip():
            return

        self.get_logger().info(
            f"Transcript [{msg.speaker_id}]: '{msg.text}'"
        )

        threading.Thread(
            target=self._generate_response,
            args=(msg.text.strip(), msg.speaker_id),
            daemon=True,
        ).start()

    # ── LLM inference ─────────────────────────────────────────────────────────

    def _generate_response(self, user_text: str, speaker_id: str) -> None:
        """Generate LLM response with streaming. Runs in thread."""
        if self._llm is None:
            self.get_logger().warn("LLM not loaded yet, dropping utterance")
            return

        with self._lock:
            if self._generating:
                self.get_logger().warn("LLM busy, dropping utterance")
                return
            self._generating = True
            self._turn_counter += 1
            turn_id = self._turn_counter

        try:
            ctx = self._ctx_store.get(speaker_id)

            # Summary compression if context is long
            if ctx.needs_compression():
                self._compress_context(ctx)

            ctx.add_user(user_text)

            prompt = build_llama_prompt(
                ctx, user_text, self._system_prompt
            )

            t0 = time.perf_counter()
            full_response = ""

            if self._stream:
                output = self._llm(
                    prompt,
                    max_tokens=self._max_tokens,
                    temperature=self._temperature,
                    top_p=self._top_p,
                    stream=True,
                    stop=["<|user|>", "<|system|>", "\n\n\n"],
                )
                for chunk in output:
                    token = chunk["choices"][0]["text"]
                    full_response += token
                    # Publish partial after each sentence boundary for low TTS latency
                    if token.endswith((".", "!", "?", "\n")):
                        self._publish_response(
                            full_response.strip(), speaker_id, turn_id, is_partial=True
                        )
            else:
                output = self._llm(
                    prompt,
                    max_tokens=self._max_tokens,
                    temperature=self._temperature,
                    top_p=self._top_p,
                    stream=False,
                    stop=["<|user|>", "<|system|>"],
                )
                full_response = output["choices"][0]["text"]

            full_response = full_response.strip()
            latency_ms = (time.perf_counter() - t0) * 1000
            self.get_logger().info(
                f"LLM [{speaker_id}] ({latency_ms:.0f}ms): '{full_response[:80]}'"
            )

            ctx.add_assistant(full_response)
            self._publish_response(full_response, speaker_id, turn_id, is_partial=False)

        except Exception as e:
            self.get_logger().error(f"LLM inference error: {e}")
        finally:
            with self._lock:
                self._generating = False

    def _compress_context(self, ctx) -> None:
        """Ask LLM to summarize old turns for context compression."""
        if self._llm is None:
            ctx.compress("(history omitted)")
            return
        try:
            summary_prompt = needs_summary_prompt(ctx)
            result = self._llm(summary_prompt, max_tokens=80, temperature=0.3, stream=False)
            summary = result["choices"][0]["text"].strip()
            ctx.compress(summary)
            self.get_logger().debug(
                f"Context compressed for {ctx.person_id}: '{summary[:60]}'"
            )
        except Exception:
            ctx.compress("(history omitted)")

    # ── Publish ───────────────────────────────────────────────────────────────

    def _publish_response(
        self, text: str, speaker_id: str, turn_id: int, is_partial: bool
    ) -> None:
        msg = ConversationResponse()
        msg.header.stamp = self.get_clock().now().to_msg()
        msg.text = text
        msg.speaker_id = speaker_id
        msg.is_partial = is_partial
        msg.turn_id = turn_id
        self._resp_pub.publish(msg)

    def _save_context(self) -> None:
        try:
            self._ctx_store.save()
        except Exception as e:
            self.get_logger().error(f"Context save error: {e}")

    def destroy_node(self) -> None:
        self._save_context()
        super().destroy_node()


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