"""记忆注入中间件 — 替代 DeerFlow MemoryMiddleware"""

import json
import logging
import os
from typing import Optional

logger = logging.getLogger(__name__)


class MemoryInjector:
    """轻量级记忆注入器 — JSON 文件持久化"""

    def __init__(self, enabled: bool = True,
                 storage_path: str = "data/memory.json"):
        self.enabled = enabled
        self.storage_path = storage_path
        self._data: dict[str, dict] = {}

        if self.enabled:
            self._load()

    def _load(self):
        """从 JSON 文件加载记忆数据"""
        if os.path.exists(self.storage_path):
            try:
                with open(self.storage_path, "r", encoding="utf-8") as f:
                    self._data = json.load(f)
                logger.info(f"Loaded {len(self._data)} memory facts")
            except Exception as e:
                logger.warning(f"Failed to load memory: {e}")
                self._data = {}

    def _save(self):
        """持久化记忆数据到 JSON 文件"""
        try:
            os.makedirs(os.path.dirname(self.storage_path), exist_ok=True)
            with open(self.storage_path, "w", encoding="utf-8") as f:
                json.dump(self._data, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.warning(f"Failed to save memory: {e}")

    def add_fact(self, key: str, value: str, confidence: float = 0.9):
        """添加记忆事实"""
        if not self.enabled:
            return
        self._data[key] = {
            "value": value,
            "confidence": confidence,
        }
        self._save()

    def get_context_string(self, max_tokens: int = 500) -> str:
        """生成用于 system prompt 注入的记忆上下文字符串"""
        if not self.enabled or not self._data:
            return ""

        lines = []
        char_budget = max_tokens * 2  # rough chars estimate
        for key, info in self._data.items():
            line = f"- {key}：{info['value']}（置信度：{info.get('confidence', 0.9)}）"
            if sum(len(l) for l in lines) + len(line) > char_budget:
                break
            lines.append(line)

        if not lines:
            return ""

        return "<memory>\n" + "\n".join(lines) + "\n</memory>"
