"""
分层记忆引擎 — coco 的核心

记忆分层架构:
  Tier 0 - 身份记忆 (identity):   永久加载，定义智能体人格
  Tier 1 - 工作记忆 (working):     当前会话状态
  Tier 2 - 近期情景 (episodic):    最近 N 次对话
  Tier 3 - 归档情景 (archived):    旧对话的压缩摘要
  Tier 4 - 语义记忆 (semantic):    提取的事实/知识
  Tier 5 - 程序记忆 (procedural):  学习到的技能/模式

长期记忆策略:
  1. 每次对话结束 → 写入情景记忆
  2. 定期触发 → 压缩旧情景为摘要
  3. 摘要中提取 → 事实写入语义记忆
  4. 检索时 → 按优先级 + 相关性混合召回
"""

import json
from pathlib import Path
from datetime import datetime, timedelta
from dataclasses import dataclass, field
from typing import Optional

from .episodic import EpisodicMemory
from .semantic import SemanticMemory
from .working import WorkingMemory
from .procedural import ProceduralMemory


@dataclass
class MemoryModule:
    """记忆模块元数据"""
    module_id: str
    memory_type: str          # identity | public | working | episodic | semantic | procedural
    path: str                 # identity 相对 agent_dir, 其余相对 memory_dir
    always_loaded: bool = False
    priority: float = 1.0
    tags: list = field(default_factory=list)
    created_at: str = ""
    updated_at: str = ""
    token_estimate: int = 0   # 估算 token 数
    summary: str = ""         # 一行摘要


class MemoryEngine:
    """
    分层记忆引擎

    用法:
        engine = MemoryEngine(agent_dir)
        engine.load_module("episodic_2026-07-04")
        engine.unload_module("episodic_2026-06-01")
        context = engine.assemble_context(token_budget=6000)
    """

    def __init__(
        self,
        agent_dir: Path,
        memory_dir: Path = None,
        llm_call=None,
        raw_retention_days: int = 30,
        summary_retention_days: int = 180,
        semantic_decay_days: int = 90,
        semantic_forget_days: int = 180,
    ):
        """
        Args:
            agent_dir: 智能体根目录 (包含 config.yaml, IDENTITY.md, public_memory/)
            memory_dir: 用户与该智能体的关系记忆目录
            llm_call: 可选, LLM 调用函数, 用于摘要生成和事实提取
        """
        self.agent_dir = Path(agent_dir)
        self.memory_dir = Path(memory_dir) if memory_dir else self.agent_dir / "memory"
        self.memory_dir.mkdir(parents=True, exist_ok=True)
        self.public_memory_dir = self.agent_dir / "public_memory"
        self.public_memory_dir.mkdir(parents=True, exist_ok=True)
        self.raw_retention_days = max(1, int(raw_retention_days or 30))
        self.summary_retention_days = max(1, int(summary_retention_days or 180))
        self.semantic_decay_days = max(1, int(semantic_decay_days or 90))
        self.semantic_forget_days = max(self.semantic_decay_days, int(semantic_forget_days or 180))

        self._llm_call = llm_call
        self._manifest_path = self.memory_dir / "manifest.json"
        self._manifest = self._load_manifest()
        self._active_ids: set = set()  # 当前加载的模块 ID

        # 初始化各层记忆处理器
        self.episodic = EpisodicMemory(self.memory_dir / "episodic")
        self.semantic = SemanticMemory(self.memory_dir / "semantic")
        self.working = WorkingMemory(self.memory_dir / "working")
        self.procedural = ProceduralMemory(self.memory_dir / "procedural")
        self.summaries_dir = self.memory_dir / "summaries"
        self.summaries_dir.mkdir(parents=True, exist_ok=True)

        # 启动时加载标识记忆 + 工作记忆 + 近期情景
        self._bootstrap()

    # ─── Manifest 管理 ───────────────────────────────────────

    def _load_manifest(self) -> dict:
        """加载记忆模块注册表"""
        if self._manifest_path.exists():
            with open(self._manifest_path, "r", encoding="utf-8") as f:
                manifest = json.load(f)
                manifest.setdefault("modules", {})
                manifest.setdefault("active", [])
                return manifest
        return {"modules": {}, "active": []}

    def _save_manifest(self):
        """保存注册表"""
        self._manifest["active"] = sorted(self._active_ids)
        with open(self._manifest_path, "w", encoding="utf-8") as f:
            json.dump(self._manifest, f, ensure_ascii=False, indent=2)

    def _register_module(self, module: MemoryModule, save: bool = True):
        """注册新模块"""
        self._manifest["modules"][module.module_id] = {
            "type": module.memory_type,
            "path": module.path,
            "always_loaded": module.always_loaded,
            "priority": module.priority,
            "tags": module.tags,
            "created_at": module.created_at,
            "updated_at": module.updated_at,
            "token_estimate": module.token_estimate,
            "summary": module.summary,
        }
        if save:
            self._save_manifest()

    # ─── 启动引导 ─────────────────────────────────────────

    def _bootstrap(self):
        """启动时加载基础记忆"""
        # 创建 identity 模块（指向 IDENTITY.md）
        identity_path = self.agent_dir / "IDENTITY.md"
        if identity_path.exists():
            identity_module = MemoryModule(
                module_id="identity",
                memory_type="identity",
                path="IDENTITY.md",
                always_loaded=True,
                priority=10.0,
                tags=["core"],
                created_at=datetime.fromtimestamp(
                    identity_path.stat().st_mtime
                ).isoformat(),
                token_estimate=self._estimate_tokens(
                    identity_path.read_text(encoding="utf-8")
                ),
                summary="智能体核心身份定义",
            )
            self._register_module(identity_module, save=False)

        self._sync_public_memory_modules()
        self._sync_episodic_modules()

        stored_active = set(self._manifest.get("active", []))
        self._active_ids = {
            mid for mid in stored_active
            if mid in self._manifest.get("modules", {})
        }

        # 加载 always_loaded 的模块
        for mid, info in self._manifest.get("modules", {}).items():
            if info.get("always_loaded"):
                self._active_ids.add(mid)

        # 首次启动时自动加载最近情景，之后尊重 manifest 中的 active 状态
        if not stored_active:
            recent = self.episodic.list_recent(limit=10)
            for ep in recent:
                self._active_ids.add(ep.module_id)

        self._save_manifest()

    def _sync_episodic_modules(self):
        """把已存在的情景记忆文件同步到 manifest。"""
        episodic_dir = self.memory_dir / "episodic"
        for fp in sorted(episodic_dir.glob("*.json")):
            module_id = f"episodic_{fp.stem}"
            existing = self._manifest["modules"].get(module_id, {})
            if module_id not in self._manifest["modules"] or str(existing.get("path", "")).startswith("memory/"):
                self._register_episodic_module(module_id, save=False)

    def _sync_public_memory_modules(self):
        """把智能体可流传的公开记忆同步到 manifest。"""
        for fp in sorted(self.public_memory_dir.glob("*")):
            if not fp.is_file() or fp.suffix.lower() not in {".md", ".txt"}:
                continue
            module_id = f"public_{fp.stem}"
            content = fp.read_text(encoding="utf-8")
            module = MemoryModule(
                module_id=module_id,
                memory_type="public",
                path=str(fp.relative_to(self.agent_dir)).replace("\\", "/"),
                always_loaded=True,
                priority=8.0,
                tags=["public", "persona"],
                created_at=datetime.fromtimestamp(fp.stat().st_mtime).isoformat(),
                updated_at=datetime.fromtimestamp(fp.stat().st_mtime).isoformat(),
                token_estimate=self._estimate_tokens(content),
                summary=content.strip().splitlines()[0][:80] if content.strip() else "公开记忆",
            )
            self._register_module(module, save=False)

    def _register_episodic_module(
        self,
        module_id: str,
        save: bool = True,
        activate: bool = False,
        touch: bool = False,
    ):
        path = self._episodic_path(module_id)
        existing = self._manifest["modules"].get(module_id, {})
        created_at = existing.get("created_at")
        if not created_at and path.exists():
            created_at = datetime.fromtimestamp(path.stat().st_mtime).isoformat()
        if not created_at:
            created_at = datetime.now().isoformat()
        updated_at = existing.get("updated_at", "")
        if touch:
            updated_at = datetime.now().isoformat()
        elif not updated_at and path.exists():
            updated_at = datetime.fromtimestamp(path.stat().st_mtime).isoformat()

        content = self.episodic.read(module_id)
        module = MemoryModule(
            module_id=module_id,
            memory_type="episodic",
            path=str(path.relative_to(self.memory_dir)).replace("\\", "/"),
            always_loaded=False,
            priority=existing.get("priority", 1.0),
            tags=existing.get("tags", []),
            created_at=created_at,
            updated_at=updated_at,
            token_estimate=self._estimate_tokens(content),
            summary=existing.get("summary", ""),
        )
        self._manifest["modules"][module.module_id] = {
            "type": module.memory_type,
            "path": module.path,
            "always_loaded": module.always_loaded,
            "priority": module.priority,
            "tags": module.tags,
            "created_at": module.created_at,
            "updated_at": module.updated_at,
            "token_estimate": module.token_estimate,
            "summary": module.summary,
            "facts_extracted_entry_count": existing.get("facts_extracted_entry_count", 0),
            "facts_extracted_at": existing.get("facts_extracted_at", ""),
        }
        if activate:
            self._active_ids.add(module_id)
        if save:
            self._save_manifest()

    def _episodic_path(self, module_id: str) -> Path:
        date_str = module_id.replace("episodic_", "")
        return self.memory_dir / "episodic" / f"{date_str}.json"

    # ─── 加载/卸载 ────────────────────────────────────────

    def load_module(self, module_id: str) -> bool:
        """
        加载指定记忆模块到活跃上下文

        Args:
            module_id: 模块 ID (如 "episodic_2026-07-04")
        Returns:
            是否成功加载
        """
        if module_id not in self._manifest["modules"]:
            return False

        self._active_ids.add(module_id)
        self._save_manifest()
        return True

    def unload_module(self, module_id: str) -> bool:
        """
        卸载指定记忆模块(不删除文件, 仅从活跃上下文中移除)

        Args:
            module_id: 模块 ID
        Returns:
            是否成功卸载
        """
        info = self._manifest["modules"].get(module_id)
        if info and info.get("always_loaded"):
            return False  # 不可卸载 always_loaded 模块

        if module_id in self._active_ids:
            self._active_ids.remove(module_id)
            self._save_manifest()
            return True
        return False

    def list_modules(self, memory_type: str = None) -> list[dict]:
        """
        列出所有记忆模块

        Args:
            memory_type: 可选, 按类型过滤
        Returns:
            模块信息列表
        """
        modules = []
        for mid, info in self._manifest.get("modules", {}).items():
            if memory_type and info.get("type") != memory_type:
                continue
            modules.append({
                "id": mid,
                "loaded": mid in self._active_ids,
                **info,
            })
        return sorted(modules, key=lambda m: m.get("updated_at", ""), reverse=True)

    def get_active_modules(self) -> list[str]:
        """获取当前活跃模块 ID 列表"""
        return sorted(self._active_ids)

    # ─── 上下文组装 ───────────────────────────────────────

    def assemble_context(
        self,
        token_budget: int = 6000,
        include_identity: bool = True,
        include_working: bool = True,
        query: str = "",
    ) -> str:
        """
        将活跃记忆组装为 LLM 上下文

        组装顺序:
          1. 身份记忆 (最高优先级)
          2. 工作记忆 (当前会话)
          3. 近期情景记忆
          4. 相关语义记忆
          5. 相关程序记忆

        Args:
            token_budget: 上下文总 token 预算
            include_identity: 是否包含身份记忆
            include_working: 是否包含工作记忆
            query: 当前用户问题, 用于召回相关语义记忆
        Returns:
            组装后的上下文字符串
        """
        sections = []
        used_tokens = 0

        def can_add(est_tokens: int) -> bool:
            return used_tokens + est_tokens <= token_budget * 0.95

        # 1. 身份记忆
        if include_identity and "identity" in self._active_ids:
            content = self._read_module_content("identity")
            est = self._estimate_tokens(content)
            if can_add(est):
                sections.append(f"<!-- 身份记忆 -->\n{content}")
                used_tokens += est

        # 1.5. 智能体公开记忆（可随智能体流传，不含用户隐私）
        public_sections = []
        for mid in sorted(self._active_ids):
            info = self._manifest["modules"].get(mid, {})
            if info.get("type") != "public":
                continue
            content = self._read_module_content(mid)
            est = self._estimate_tokens(content)
            if can_add(est):
                public_sections.append(content)
                used_tokens += est
        if public_sections:
            sections.append("<!-- 智能体公开记忆 -->\n" + "\n---\n".join(public_sections))

        # 2. 工作记忆
        if include_working:
            wm = self.working.read()
            if wm:
                est = self._estimate_tokens(wm)
                if can_add(est):
                    sections.append(f"<!-- 工作记忆(当前会话) -->\n{wm}")
                    used_tokens += est

        # 3. 情景记忆 (按时间倒序)
        episodic_sections = []
        for mid in sorted(self._active_ids, reverse=True):
            info = self._manifest["modules"].get(mid, {})
            if info.get("type") != "episodic":
                continue
            content = self._read_module_content(mid)
            est = self._estimate_tokens(content)
            if can_add(est):
                episodic_sections.append(content)
                used_tokens += est

        if episodic_sections:
            sections.append("<!-- 情景记忆 -->\n" + "\n---\n".join(episodic_sections))

        # 4. 语义记忆 (按相关性 + 优先级)
        semantic_content = self.semantic.retrieve_relevant(
            query=query, limit=20
        )
        if semantic_content:
            text = "\n".join(semantic_content)
            est = self._estimate_tokens(text)
            if can_add(est):
                sections.append(f"<!-- 语义记忆(事实/知识) -->\n{text}")
                used_tokens += est

        # 5. 程序记忆
        proc_content = self.procedural.retrieve_relevant(limit=5)
        if proc_content:
            text = "\n".join(proc_content)
            est = self._estimate_tokens(text)
            if can_add(est):
                sections.append(f"<!-- 程序记忆(技能/模式) -->\n{text}")
                used_tokens += est

        return "\n\n".join(sections)

    def _read_module_content(self, module_id: str) -> str:
        """读取模块内容"""
        info = self._manifest["modules"].get(module_id)
        if not info:
            return ""

        path = self._module_path(info)
        if not path.exists():
            return ""

        # 情景记忆走专门的处理逻辑
        if info.get("type") == "episodic":
            return self.episodic.read(module_id)

        return path.read_text(encoding="utf-8")

    def _module_path(self, info: dict) -> Path:
        path = Path(info.get("path", ""))
        if path.is_absolute():
            return path
        if info.get("type") in {"identity", "public"}:
            return self.agent_dir / path
        return self.memory_dir / path

    # ─── 记忆写入 ────────────────────────────────────────

    def record_conversation(
        self, user_msg: str, assistant_msg: str, metadata: dict = None
    ):
        """
        记录一轮对话

        Args:
            user_msg: 用户消息
            assistant_msg: 助手回复
            metadata: 额外元数据
        """
        today = datetime.now().strftime("%Y-%m-%d")
        module_id = f"episodic_{today}"
        self._register_episodic_module(module_id, save=False, activate=True)

        self.episodic.append(module_id, {
            "role": "user",
            "content": user_msg,
            "timestamp": datetime.now().isoformat(),
        })
        self.episodic.append(module_id, {
            "role": "assistant",
            "content": assistant_msg,
            "timestamp": datetime.now().isoformat(),
        })

        # 更新工作记忆
        self.working.update_last_interaction(datetime.now().isoformat())
        self._register_episodic_module(module_id, save=True, activate=True, touch=True)

    def end_session(self):
        """
        会话结束时的记忆整理:
        1. 压缩今日情景记忆, 生成摘要
        2. 提取关键事实到语义记忆
        3. 检查是否需要归档旧记忆
        """
        today = datetime.now().strftime("%Y-%m-%d")
        module_id = f"episodic_{today}"

        if not self._episodic_path(module_id).exists():
            self._check_archive()
            self.working.clear()
            return

        self._register_episodic_module(module_id, save=False, activate=True)

        # 生成摘要
        summary = self.episodic.summarize(module_id, llm_call=self._llm_call)
        if summary:
            self._write_summary(module_id, summary)
            # 更新模块元数据
            if module_id in self._manifest["modules"]:
                self._manifest["modules"][module_id]["summary"] = summary
                self._manifest["modules"][module_id]["updated_at"] = \
                    datetime.now().isoformat()
                self._save_manifest()

        # 只从新增对话中提取事实，避免 close/__del__ 重复追加
        entries = self.episodic.read_raw(module_id)
        info = self._manifest["modules"].get(module_id, {})
        processed_count = int(info.get("facts_extracted_entry_count", 0) or 0)
        new_entries = entries[processed_count:]
        content = self._format_episodic_entries(new_entries)
        if content and self._llm_call:
            facts = self._extract_facts(content)
            for fact in facts:
                self.semantic.add(fact, tags=["auto"], source=module_id)
            self._manifest["modules"][module_id]["facts_extracted_entry_count"] = len(entries)
            self._manifest["modules"][module_id]["facts_extracted_at"] = datetime.now().isoformat()
            self._save_manifest()

        # 检查归档
        self._check_archive()
        self.semantic.decay_unused(
            decay_days=self.semantic_decay_days,
            forget_days=self.semantic_forget_days,
        )

        # 清除工作记忆
        self.working.clear()

    def _write_summary(self, module_id: str, summary: str):
        date_str = module_id.replace("episodic_", "")
        path = self.summaries_dir / f"{date_str}.json"
        payload = {
            "date": date_str,
            "module_id": module_id,
            "summary": summary,
            "scope": "relationship",
            "updated_at": datetime.now().isoformat(),
        }
        with open(path, "w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False, indent=2)

    def _format_episodic_entries(self, entries: list[dict]) -> str:
        lines = []
        for entry in entries:
            role = "用户" if entry.get("role") == "user" else "助手"
            content = entry.get("content", "")
            if content:
                lines.append(f"[{role}] {content}")
        return "\n".join(lines)

    def _extract_facts(self, conversation_text: str) -> list[str]:
        """从对话中提取关键事实"""
        if not self._llm_call:
            return []

        prompt = f"""从以下对话中提取 3-5 条关键事实。每条一行, 用中文。

规则:
- 只提取对智能体有用的信息(用户偏好、项目信息、重要决定)
- 不要提取琐碎的闲聊内容
- 用第三人称陈述(如 "公子喜欢用Python")

对话:
{conversation_text[:3000]}

关键事实:"""

        try:
            response = self._llm_call(prompt, max_tokens=500)
            facts = [
                line.strip("- ").strip()
                for line in response.split("\n")
                if line.strip() and len(line.strip()) > 10
            ]
            return facts[:5]
        except Exception:
            return []

    def _check_archive(self):
        """
        检查并归档旧记忆:
        超过阈值的旧情景记忆 → 生成超压缩摘要 → 原始对话可归档
        """
        threshold = datetime.now() - timedelta(days=self.raw_retention_days)
        for mid in list(self._active_ids):
            if not mid.startswith("episodic_"):
                continue
            info = self._manifest["modules"].get(mid, {})
            created = info.get("created_at", "")
            if created and datetime.fromisoformat(created) < threshold:
                # 从活跃列表中移除(但保留文件)
                if not info.get("always_loaded"):
                    self._active_ids.discard(mid)

        summary_threshold = datetime.now() - timedelta(days=self.summary_retention_days)
        for fp in self.summaries_dir.glob("*.json"):
            try:
                summary_date = datetime.fromisoformat(fp.stem)
            except ValueError:
                continue
            if summary_date < summary_threshold:
                fp.unlink(missing_ok=True)

        self._save_manifest()

    # ─── 记忆检索 ────────────────────────────────────────

    def search(self, query: str, memory_type: str = None, limit: int = 5) -> list[dict]:
        """
        搜索记忆(关键词匹配)

        Args:
            query: 搜索关键词
            memory_type: 可选, 限定记忆类型
            limit: 返回条数
        Returns:
            匹配的记忆条目
        """
        results = []

        # 搜索情景记忆
        if not memory_type or memory_type == "episodic":
            results.extend(self.episodic.search(query, limit=limit))

        # 搜索语义记忆
        if not memory_type or memory_type == "semantic":
            results.extend(self.semantic.search(query, limit=limit))

        # 搜索程序记忆
        if not memory_type or memory_type == "procedural":
            results.extend(self.procedural.search(query, limit=limit))

        # 按相关度排序
        results.sort(key=lambda r: r.get("score", 0), reverse=True)
        return results[:limit]

    # ─── 辅助方法 ────────────────────────────────────────

    def _estimate_tokens(self, text: str) -> int:
        """粗略估算 token 数 (中文约 1.5 字符/token, 英文约 4 字符/token)"""
        if not text:
            return 0
        chinese_chars = sum(1 for c in text if "\u4e00" <= c <= "\u9fff")
        other_chars = len(text) - chinese_chars
        return int(chinese_chars / 1.5 + other_chars / 4)

    def get_stats(self) -> dict:
        """获取记忆统计信息"""
        episodic_count = self.episodic.count()
        semantic_count = self.semantic.count()
        procedural_count = self.procedural.count()
        summary_count = len(list(self.summaries_dir.glob("*.json")))

        total_tokens = sum(
            self._manifest["modules"].get(mid, {}).get("token_estimate", 0)
            for mid in self._active_ids
        )

        return {
            "active_modules": len(self._active_ids),
            "total_modules": len(self._manifest.get("modules", {})),
            "episodic_entries": episodic_count,
            "summaries": summary_count,
            "semantic_facts": semantic_count,
            "procedural_skills": procedural_count,
            "estimated_context_tokens": total_tokens,
            "memory_dir": str(self.memory_dir),
        }
