"""conversation_node.py — Local LLM conversation engine with per-person context.
Issue #83/#161/#167
"""
from __future__ import annotations
import json, threading, time
from typing import Dict, 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 SpeechTranscript, ConversationResponse
from .llm_context import ContextStore, build_llama_prompt, load_system_prompt, needs_summary_prompt

_LANG_NAMES: Dict[str, str] = {
    "en": "English", "fr": "French", "es": "Spanish", "de": "German",
    "it": "Italian", "pt": "Portuguese", "ja": "Japanese", "zh": "Chinese",
    "ko": "Korean", "ar": "Arabic", "ru": "Russian", "nl": "Dutch",
    "pl": "Polish", "sv": "Swedish", "da": "Danish", "fi": "Finnish",
    "no": "Norwegian", "tr": "Turkish", "hi": "Hindi", "uk": "Ukrainian",
    "cs": "Czech", "ro": "Romanian", "hu": "Hungarian", "el": "Greek",
}

class ConversationNode(Node):
    """Local LLM inference node with per-person conversation memory."""
    def __init__(self) -> None:
        super().__init__("conversation_node")
        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
        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)
        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()
        self._speaker_lang: Dict[str, str] = {}
        self._emotions: Dict[str, str] = {}
        self.create_subscription(String, "/social/emotion/context", self._on_emotion_context, 10)
        threading.Thread(target=self._load_llm, daemon=True).start()
        self._save_timer = self.create_timer(self._save_interval, self._save_context)
        self.get_logger().info(
            f"ConversationNode init (model={self._model_path}, "
            f"gpu_layers={self._n_gpu}, ctx={self._n_ctx})"
        )

    def _load_llm(self) -> None:
        t0 = time.time()
        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)")
        except Exception as e:
            self.get_logger().error(f"LLM load failed: {e}")

    def _on_transcript(self, msg: SpeechTranscript) -> None:
        if msg.is_partial or not msg.text.strip():
            return
        if msg.language:
            self._speaker_lang[msg.speaker_id] = msg.language
        self.get_logger().info(f"Transcript [{msg.speaker_id}/{msg.language or '?'}]: '{msg.text}'")
        threading.Thread(target=self._generate_response, args=(msg.text.strip(), msg.speaker_id), daemon=True).start()

    def _generate_response(self, user_text: str, speaker_id: str) -> None:
        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
        lang = self._speaker_lang.get(speaker_id, "en")
        try:
            ctx = self._ctx_store.get(speaker_id)
            if ctx.needs_compression():
                self._compress_context(ctx)
            emotion_hint = self._emotion_hint(speaker_id)
            lang_hint = self._language_hint(speaker_id)
            hints = " ".join(h for h in (emotion_hint, lang_hint) if h)
            annotated = f"{user_text}  {hints}".rstrip() if hints else user_text
            ctx.add_user(annotated)
            prompt = build_llama_prompt(ctx, annotated, 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
                    if token.endswith((".", "!", "?", "\n")):
                        self._publish_response(full_response.strip(), speaker_id, turn_id, language=lang, 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()
            self.get_logger().info(f"LLM [{speaker_id}/{lang}] ({(time.perf_counter()-t0)*1000:.0f}ms): '{full_response[:80]}'")
            ctx.add_assistant(full_response)
            self._publish_response(full_response, speaker_id, turn_id, language=lang, 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:
        if self._llm is None: ctx.compress("(history omitted)"); return
        try:
            result = self._llm(needs_summary_prompt(ctx), max_tokens=80, temperature=0.3, stream=False)
            ctx.compress(result["choices"][0]["text"].strip())
        except Exception: ctx.compress("(history omitted)")

    def _language_hint(self, speaker_id: str) -> str:
        lang = self._speaker_lang.get(speaker_id, "en")
        if lang and lang != "en":
            return f"[Please respond in {_LANG_NAMES.get(lang, lang)}.]"
        return ""

    def _on_emotion_context(self, msg: String) -> None:
        try:
            for k, v in json.loads(msg.data).get("emotions", {}).items():
                self._emotions[k] = v
        except Exception: pass

    def _emotion_hint(self, speaker_id: str) -> str:
        emo = self._emotions.get(speaker_id, "")
        return f"[The person seems {emo} right now.]" if emo and emo != "neutral" else ""

    def _publish_response(self, text: str, speaker_id: str, turn_id: int, language: str = "en", is_partial: bool = False) -> 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; msg.language = language
        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()
