"""
情景记忆 — 对话记录与检索
存储格式: memory/episodic/YYYY-MM-DD.json
"""

import json
import re
from pathlib import Path
from datetime import datetime
from typing import Optional


class EpisodicMemory:
    """情景记忆处理器: 按日期存储对话记录"""

    def __init__(self, storage_dir: Path):
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)

    def _file_path(self, module_id: str) -> Path:
        """模块 ID 转文件路径, 如 episodic_2026-07-04 → .../2026-07-04.json"""
        date_str = module_id.replace("episodic_", "")
        return self.storage_dir / f"{date_str}.json"

    def append(self, module_id: str, entry: dict):
        """追加一条对话记录"""
        path = self._file_path(module_id)
        entries = self._read_raw(path)
        entries.append(entry)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(entries, f, ensure_ascii=False, indent=2)

    def read(self, module_id: str) -> str:
        """读取某天的对话为可读文本"""
        path = self._file_path(module_id)
        entries = self._read_raw(path)
        if not entries:
            return ""

        lines = []
        for entry in entries:
            role = "用户" if entry["role"] == "user" else "助手"
            lines.append(f"[{role}] {entry['content']}")
        return "\n".join(lines)

    def read_raw(self, module_id: str) -> list[dict]:
        """读取原始 JSON 条目"""
        return self._read_raw(self._file_path(module_id))

    def summarize(self, module_id: str, llm_call=None) -> Optional[str]:
        """
        生成对话摘要

        Args:
            module_id: 模块 ID
            llm_call: LLM 调用函数
        Returns:
            摘要文本, 如果无法生成则返回 None
        """
        content = self.read(module_id)
        if not content or not llm_call:
            return None

        # 如果内容太少, 不需要摘要
        if len(content) < 200:
            return content[:200]

        prompt = f"""请用 2-3 句话总结以下对话的主要内容。只输出摘要, 不要其他内容。

{content[:4000]}

摘要:"""

        try:
            return llm_call(prompt, max_tokens=200).strip()
        except Exception:
            return None

    def list_recent(self, limit: int = 10) -> list:
        """列出最近的 N 天记忆"""
        from dataclasses import dataclass, field

        @dataclass
        class Module:
            module_id: str

        files = sorted(self.storage_dir.glob("*.json"), reverse=True)
        modules = []
        for fp in files[:limit]:
            date_str = fp.stem
            module_id = f"episodic_{date_str}"
            entries = self._read_raw(fp)
            modules.append(type("Module", (), {
                "module_id": module_id,
                "entry_count": len(entries),
                "path": str(fp),
            })())
        return modules

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

        Args:
            query: 搜索词
            limit: 返回条数
        Returns:
            [{score, content, date, module_id}, ...]
        """
        results = []
        query_lower = query.lower()

        for fp in sorted(self.storage_dir.glob("*.json"), reverse=True):
            entries = self._read_raw(fp)
            for entry in entries:
                content = entry.get("content", "")
                if query_lower in content.lower():
                    # 简单 TF 得分
                    score = content.lower().count(query_lower)
                    results.append({
                        "score": score,
                        "date": fp.stem,
                        "module_id": f"episodic_{fp.stem}",
                        "role": entry.get("role", ""),
                        "content": content[:500],
                        "timestamp": entry.get("timestamp", ""),
                    })

        results.sort(key=lambda r: r["score"], reverse=True)
        return results[:limit]

    def count(self) -> int:
        """总条目数"""
        total = 0
        for fp in self.storage_dir.glob("*.json"):
            total += len(self._read_raw(fp))
        return total

    def _read_raw(self, path: Path) -> list[dict]:
        """读取原始 JSON"""
        if not path.exists():
            return []
        try:
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
        except (json.JSONDecodeError, FileNotFoundError):
            return []
