#!/usr/bin/env python3
"""
agint_memory.memory_node - ROS2 Memory Node

ROS2 节点，负责记录对话历史，提供短记忆给上游节点使用。

Topic 设计：
    订阅：
        - /asr (std_msgs/String) - 用户语音识别文本
        - /llm_response (std_msgs/String) - OmniNode 流式响应（需聚合）
        - /plan (std_msgs/String) - 计划 JSON（可选）
    发布：
        - /memory/context (std_msgs/String) - 记忆上下文 JSON

Author: Jonathan
"""

import json
import time
import threading
from typing import List, Dict

import rclpy
from rclpy.node import Node
from rclpy.qos import QoSProfile, ReliabilityPolicy, HistoryPolicy
from std_msgs.msg import String
from std_srvs.srv import Empty


class MemoryNode(Node):
    """ROS2 Memory Node - 对话记忆管理

    核心职责：
        1. 订阅 /asr，记录用户输入
        2. 订阅 /llm_response，聚合流式回复
        3. 保存 raw memory，控制 2MB
        4. 生成 short_memory（最近 15 行）
        5. 发布 /memory/context
    """

    def __init__(self) -> None:
        """节点初始化"""
        super().__init__("memory_node")
        self.get_logger().info("Memory Node is starting...")

        # =====================================================================
        # 1. 声明并读取 ROS2 参数
        # =====================================================================
        self.declare_parameter("input_asr_topic", "/asr")
        self.declare_parameter("input_llm_response_topic", "/llm_response")
        self.declare_parameter("input_plan_topic", "/plan")
        self.declare_parameter("output_memory_context_topic", "/memory/context")

        self.declare_parameter("session_id", "default")
        self.declare_parameter("max_raw_memory_bytes", 2 * 1024 * 1024)
        self.declare_parameter("short_memory_lines", 15)
        self.declare_parameter("assistant_flush_timeout_sec", 1.0)
        self.declare_parameter("publish_on_update", True)
        self.declare_parameter("debug_log", True)

        self.input_asr_topic = self.get_parameter("input_asr_topic").value
        self.input_llm_response_topic = self.get_parameter("input_llm_response_topic").value
        self.input_plan_topic = self.get_parameter("input_plan_topic").value
        self.output_memory_context_topic = self.get_parameter("output_memory_context_topic").value

        self.session_id = self.get_parameter("session_id").value
        self.max_raw_memory_bytes = self.get_parameter("max_raw_memory_bytes").value
        self.short_memory_lines = self.get_parameter("short_memory_lines").value
        self.assistant_flush_timeout_sec = self.get_parameter("assistant_flush_timeout_sec").value
        self.publish_on_update = self.get_parameter("publish_on_update").value
        self.debug_log = self.get_parameter("debug_log").value

        # =====================================================================
        # 2. 初始化 QoS 配置
        # =====================================================================
        qos = QoSProfile(
            reliability=ReliabilityPolicy.RELIABLE,
            history=HistoryPolicy.KEEP_LAST,
            depth=10,
        )

        # =====================================================================
        # 3. 创建 Publisher
        # =====================================================================
        self.memory_pub = self.create_publisher(
            String,
            self.output_memory_context_topic,
            qos
        )

        # =====================================================================
        # 4. 创建 Subscriber
        # =====================================================================
        self.asr_sub = self.create_subscription(
            String,
            self.input_asr_topic,
            self.on_asr,
            qos,
        )

        self.llm_response_sub = self.create_subscription(
            String,
            self.input_llm_response_topic,
            self.on_llm_response,
            qos,
        )

        self.plan_sub = self.create_subscription(
            String,
            self.input_plan_topic,
            self.on_plan,
            qos,
        )

        # =====================================================================
        # 5. 创建 Timer（用于 flush 检查）
        # =====================================================================
        self.flush_timer = self.create_timer(0.2, self._check_flush_timeout)

        # =====================================================================
        # 6. 创建 Service（用于清除记忆）
        # =====================================================================
        self.clear_memory_srv = self.create_service(
            Empty,
            "/memory/clear",
            self._on_clear_memory
        )

        # =====================================================================
        # 7. 内部数据结构
        # =====================================================================
        self.turns: List[Dict] = []          # 所有对话轮次
        self.long_memory: str = ""          # 长期记忆（V1 为空）
        self.assistant_buffer: str = ""      # 流式回复缓冲区
        self.assistant_last_update: float = 0.0

        self.turn_counter: int = 0           # 总轮次计数
        self._lock = threading.Lock()        # 线程锁

        # =====================================================================
        # 7. 发布初始记忆
        # =====================================================================
        self.publish_context()

        self.get_logger().info("Memory Node initialized successfully")
        self.get_logger().info(f"  - ASR Topic: {self.input_asr_topic}")
        self.get_logger().info(f"  - LLM Response Topic: {self.input_llm_response_topic}")
        self.get_logger().info(f"  - Plan Topic: {self.input_plan_topic}")
        self.get_logger().info(f"  - Memory Context Topic: {self.output_memory_context_topic}")
        self.get_logger().info(f"  - Session ID: {self.session_id}")
        self.get_logger().info(f"  - Max Memory Size: {self.max_raw_memory_bytes} bytes (~{self.max_raw_memory_bytes // (1024 * 1024)} MB)")
        self.get_logger().info(f"  - Short Memory Lines: {self.short_memory_lines}")

    # =========================================================================
    # 订阅回调 - ASR
    # =========================================================================
    def on_asr(self, msg: String) -> None:
        """处理 ASR 语音识别文本"""
        text = msg.data.strip()
        if not text:
            return

        with self._lock:
            self.append_turn("user", text)
            self._maybe_trim_memory()
            self.publish_context_if_needed()

        if self.debug_log:
            self.get_logger().debug(f"[MemoryNode] user: {text[:50]}...")

    # =========================================================================
    # 订阅回调 - LLM Response（流式）
    # =========================================================================
    def on_llm_response(self, msg: String) -> None:
        """处理 OmniNode 流式响应

        流式片段累加到缓冲区，超时或收到空字符串时 flush。
        """
        text = msg.data

        with self._lock:
            if text and text.strip():
                self.assistant_buffer += text
                self.assistant_last_update = time.time()
                if self.debug_log:
                    self.get_logger().debug(f"[MemoryNode] llm_response chunk: {text[:30]}...")
            else:
                self.flush_assistant_buffer("empty_message")
                self.publish_context_if_needed()

    # =========================================================================
    # 订阅回调 - Plan（可选）
    # =========================================================================
    def on_plan(self, msg: String) -> None:
        """处理 Plan 消息（可选，V1 暂不处理）"""
        pass

    # =========================================================================
    # Timer 回调 - 检查 flush 超时
    # =========================================================================
    def _check_flush_timeout(self) -> None:
        """定时检查是否需要 flush assistant buffer"""
        with self._lock:
            if self.assistant_buffer:
                elapsed = time.time() - self.assistant_last_update
                if elapsed >= self.assistant_flush_timeout_sec:
                    self.flush_assistant_buffer("timeout")
                    self.publish_context_if_needed()

    # =========================================================================
    # 内部方法
    # =========================================================================
    def append_turn(self, role: str, text: str) -> None:
        """添加一轮对话"""
        text = text.strip()
        if not text:
            return

        if self.turns and self.turns[-1]["role"] == role and self.turns[-1]["text"] == text:
            if self.debug_log:
                self.get_logger().debug(f"[MemoryNode] Duplicate turn skipped")
            return

        turn = {
            "role": role,
            "text": text,
            "timestamp": time.time(),
        }
        self.turns.append(turn)
        self.turn_counter += 1

        if self.debug_log:
            self.get_logger().debug(f"[MemoryNode] Added turn [{role}]: {text[:50]}... (total: {len(self.turns)})")

    def flush_assistant_buffer(self, reason: str) -> None:
        """将 assistant buffer 写入 turns"""
        text = self.assistant_buffer.strip()
        self.assistant_buffer = ""

        if not text:
            return

        text = self.clean_assistant_text(text)
        if not text:
            return

        self.append_turn("assistant", text)
        self.get_logger().debug(f"[MemoryNode] Flushed buffer (reason={reason})")

    def clean_assistant_text(self, text: str) -> str:
        """清洗 assistant 回复文本

        规则：
            1. 去掉 <FINAL_JSON> 及其后面的内容
            2. 去掉残留的 <FINAL_JSON 或 FINAL_JSON
            3. 纯 JSON 不保存
        """
        if "<FINAL_JSON>" in text:
            text = text.split("<FINAL_JSON>", 1)[0]

        text = text.replace("<FINAL_JSON", "")
        text = text.replace("FINAL_JSON", "")

        stripped = text.strip()
        if stripped.startswith("{"):
            try:
                json.loads(stripped)
                if self.debug_log:
                    self.get_logger().debug("[MemoryNode] Skipped JSON-only text")
                return ""
            except json.JSONDecodeError:
                pass

        return text.strip()

    def estimate_memory_bytes(self) -> int:
        """估算当前 turns 的内存占用"""
        try:
            return len(json.dumps(self.turns, ensure_ascii=False).encode("utf-8"))
        except Exception:
            return 0

    def _maybe_trim_memory(self) -> None:
        """如果内存超过上限，从最早开始删除"""
        while self.estimate_memory_bytes() > self.max_raw_memory_bytes and self.turns:
            removed = self.turns.pop(0)
            if self.debug_log:
                self.get_logger().debug(f"[MemoryNode] Trimmed: {removed['role']}: {removed['text'][:30]}...")

    def role_label(self, role: str) -> str:
        """角色映射为中文标签"""
        labels = {
            "user": "用户",
            "assistant": "助手",
            "plan": "计划",
            "system": "系统",
        }
        return labels.get(role, role)

    def build_short_memory(self) -> str:
        """构建 short_memory 文本"""
        if not self.turns:
            return "最近对话：无"

        recent = self.turns[-self.short_memory_lines:]
        lines = ["最近对话："]
        for t in recent:
            role = self.role_label(t["role"])
            text = t["text"].strip()
            if text:
                lines.append(f"{role}：{text}")

        return "\n".join(lines)

    def build_memory_text(self) -> str:
        """构建 memory_text（发给 LLM 的文本）"""
        parts = []

        if self.long_memory.strip():
            parts.append("长期记忆：\n" + self.long_memory.strip())
        else:
            parts.append("长期记忆：无")

        parts.append(self.build_short_memory())

        return "\n".join(parts)

    def publish_context_if_needed(self) -> None:
        """如果配置为每次更新后发布，则发布"""
        if self.publish_on_update:
            self.publish_context()

    def publish_context(self) -> None:
        """发布 /memory/context"""
        msg = String()
        payload = {
            "session_id": self.session_id,
            "short_memory": self.build_short_memory(),
            "long_memory": self.long_memory,
            "memory_text": self.build_memory_text(),
            "turn_count": len(self.turns),
            "total_turns": self.turn_counter,
            "updated_at": time.time(),
        }
        msg.data = json.dumps(payload, ensure_ascii=False)
        self.memory_pub.publish(msg)

        if self.debug_log:
            self.get_logger().info(
                f"[MemoryNode] Published: turns={len(self.turns)}, "
                f"total={self.turn_counter}, text_len={len(payload['memory_text'])}"
            )

    def get_memory_text(self) -> str:
        """获取当前 memory_text（供外部调用）"""
        with self._lock:
            return self.build_memory_text()

    def clear_memory(self) -> None:
        """清空所有记忆"""
        with self._lock:
            self.turns.clear()
            self.long_memory = ""
            self.assistant_buffer = ""
            self.turn_counter = 0
            self.publish_context()
            self.get_logger().info("[MemoryNode] Memory cleared")

    def _on_clear_memory(self, request, response) -> None:
        """Service 回调：清除记忆"""
        self.clear_memory()
        return response


def main(args=None) -> None:
    """节点主入口"""
    rclpy.init(args=args)
    try:
        node = MemoryNode()
        try:
            rclpy.spin(node)
        except KeyboardInterrupt:
            node.get_logger().info("KeyboardInterrupt received, shutting down.")
        finally:
            node.destroy_node()
    finally:
        rclpy.shutdown()


if __name__ == "__main__":
    main()

