"""
记忆整合器 - 批量整理对话历史

实现用户的想法:
1. 保存一整天的对话上下文
2. 空闲时段 (如凌晨) 自动整理
3. 归纳精华存入 MEMORY.md

参考:
- Claude-Mem Worker Service
- LangMem Background Manager
"""

import json
import logging
from datetime import datetime, timedelta
from pathlib import Path

from .extractor import MemoryExtractor
from .types import ConversationTurn, Memory, SessionSummary

logger = logging.getLogger(__name__)


class MemoryConsolidator:
    """记忆整合器 - 批量处理对话历史"""

    def __init__(
        self,
        data_dir: Path,
        brain=None,
        extractor: MemoryExtractor | None = None,
    ):
        """
        Args:
            data_dir: 数据目录 (存放对话历史)
            brain: LLM 大脑实例
            extractor: 记忆提取器
        """
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(parents=True, exist_ok=True)

        self.brain = brain
        self.extractor = extractor or MemoryExtractor(brain)

        # 对话历史存储目录
        self.history_dir = self.data_dir / "conversation_history"
        self.history_dir.mkdir(exist_ok=True)

        # 已整理的会话
        self.summaries_file = self.data_dir / "session_summaries.json"

    def save_conversation_turn(
        self,
        session_id: str,
        turn: ConversationTurn,
    ) -> None:
        """
        保存对话轮次 (实时保存)

        每个会话一个文件，追加写入
        """
        session_file = self.history_dir / f"{session_id}.jsonl"

        with open(session_file, "a", encoding="utf-8") as f:
            f.write(json.dumps(turn.to_dict(), ensure_ascii=False) + "\n")

    def load_session_history(self, session_id: str) -> list[ConversationTurn]:
        """加载会话历史"""
        session_file = self.history_dir / f"{session_id}.jsonl"

        if not session_file.exists():
            return []

        turns = []
        with open(session_file, encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    data = json.loads(line)
                    turn = ConversationTurn(
                        role=data["role"],
                        content=data["content"],
                        timestamp=datetime.fromisoformat(data["timestamp"]),
                        tool_calls=data.get("tool_calls", []),
                        tool_results=data.get("tool_results", []),
                    )
                    turns.append(turn)

        return turns

    def get_today_sessions(self) -> list[str]:
        """获取今天的所有会话 ID"""
        today = datetime.now().date()
        sessions = []

        for file in self.history_dir.glob("*.jsonl"):
            # 检查文件修改时间
            mtime = datetime.fromtimestamp(file.stat().st_mtime)
            if mtime.date() == today:
                sessions.append(file.stem)

        return sessions

    def get_unprocessed_sessions(self) -> list[str]:
        """获取未处理的会话"""
        # 加载已处理的会话
        processed = set()
        if self.summaries_file.exists():
            with open(self.summaries_file, encoding="utf-8") as f:
                for line in f:
                    if line.strip():
                        summary = json.loads(line)
                        processed.add(summary["session_id"])

        # 找出未处理的
        unprocessed = []
        for file in self.history_dir.glob("*.jsonl"):
            if file.stem not in processed:
                unprocessed.append(file.stem)

        return unprocessed

    async def consolidate_session(
        self,
        session_id: str,
    ) -> tuple[SessionSummary, list[Memory]]:
        """
        整理单个会话

        1. 加载对话历史
        2. 生成会话摘要
        3. 提取记忆
        """
        turns = self.load_session_history(session_id)

        if not turns:
            return None, []

        # 生成会话摘要
        summary = await self._generate_summary(session_id, turns)

        # 提取记忆
        memories = []

        # 基于规则提取
        for turn in turns:
            extracted = self.extractor.extract_from_turn(turn)
            memories.extend(extracted)

        # 使用 LLM 高级提取
        if self.brain:
            llm_memories = await self.extractor.extract_with_llm(
                turns, context=f"会话摘要: {summary.task_description}"
            )
            memories.extend(llm_memories)

        # 去重
        memories = self.extractor.deduplicate(memories, [])

        # 更新摘要中的记忆 ID
        summary.memories_created = [m.id for m in memories]

        # 保存摘要
        self._save_summary(summary)

        return summary, memories

    async def consolidate_all_unprocessed(self) -> tuple[list[SessionSummary], list[Memory]]:
        """
        整理所有未处理的会话

        适合在空闲时段 (如凌晨) 批量执行
        """
        unprocessed = self.get_unprocessed_sessions()

        all_summaries = []
        all_memories = []

        for session_id in unprocessed:
            try:
                summary, memories = await self.consolidate_session(session_id)
                if summary:
                    all_summaries.append(summary)
                    all_memories.extend(memories)
                    logger.info(f"Consolidated session {session_id}: {len(memories)} memories")
            except Exception as e:
                logger.error(f"Failed to consolidate session {session_id}: {e}")

        return all_summaries, all_memories

    async def _generate_summary(
        self,
        session_id: str,
        turns: list[ConversationTurn],
    ) -> SessionSummary:
        """使用 LLM 生成会话摘要"""

        start_time = turns[0].timestamp if turns else datetime.now()
        end_time = turns[-1].timestamp if turns else datetime.now()

        # 简单摘要 (不用 LLM)
        if not self.brain or len(turns) < 3:
            # 从用户消息提取任务描述
            user_messages = [t.content for t in turns if t.role == "user"]
            task_desc = user_messages[0][:200] if user_messages else "Unknown task"

            return SessionSummary(
                session_id=session_id,
                start_time=start_time,
                end_time=end_time,
                task_description=task_desc,
                outcome="completed",
            )

        # 使用 LLM 生成详细摘要
        conv_text = "\n".join(
            [
                f"[{turn.role}]: {turn.content[:300]}"
                for turn in turns[-30:]  # 最近30轮
            ]
        )

        prompt = f"""总结以下对话会话:

{conv_text}

请提供:
1. task_description: 用户的主要任务是什么 (一句话)
2. outcome: 任务结果 (success/partial/failed)
3. key_actions: 关键操作 (最多5个)
4. learnings: 值得记住的经验 (最多3个)
5. errors: 遇到的错误 (如果有)

用 JSON 格式输出。
"""

        try:
            response = await self.brain.think(
                prompt, system="你是一个会话分析专家，擅长提取关键信息。只输出 JSON，不要其他内容。"
            )

            # 解析 JSON
            import re

            json_match = re.search(r"\{.*\}", response, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group())
                return SessionSummary(
                    session_id=session_id,
                    start_time=start_time,
                    end_time=end_time,
                    task_description=data.get("task_description", ""),
                    outcome=data.get("outcome", "completed"),
                    key_actions=data.get("key_actions", []),
                    learnings=data.get("learnings", []),
                    errors_encountered=data.get("errors", []),
                )
        except Exception as e:
            logger.error(f"LLM summary generation failed: {e}")

        # 回退到简单摘要
        user_messages = [t.content for t in turns if t.role == "user"]
        return SessionSummary(
            session_id=session_id,
            start_time=start_time,
            end_time=end_time,
            task_description=user_messages[0][:200] if user_messages else "Unknown",
            outcome="completed",
        )

    def _save_summary(self, summary: SessionSummary) -> None:
        """保存会话摘要"""
        with open(self.summaries_file, "a", encoding="utf-8") as f:
            f.write(json.dumps(summary.to_dict(), ensure_ascii=False) + "\n")

    def get_recent_summaries(self, days: int = 7) -> list[SessionSummary]:
        """获取最近N天的会话摘要"""
        if not self.summaries_file.exists():
            return []

        cutoff = datetime.now() - timedelta(days=days)
        summaries = []

        with open(self.summaries_file, encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    data = json.loads(line)
                    end_time = datetime.fromisoformat(data["end_time"])
                    if end_time > cutoff:
                        summaries.append(
                            SessionSummary(
                                session_id=data["session_id"],
                                start_time=datetime.fromisoformat(data["start_time"]),
                                end_time=end_time,
                                task_description=data.get("task_description", ""),
                                outcome=data.get("outcome", ""),
                                key_actions=data.get("key_actions", []),
                                learnings=data.get("learnings", []),
                                errors_encountered=data.get("errors_encountered", []),
                                memories_created=data.get("memories_created", []),
                            )
                        )

        return summaries

    def cleanup_old_history(self, days: int = 30) -> int:
        """
        清理旧的对话历史文件（按天数）

        保留摘要和记忆，删除原始对话
        """
        cutoff = datetime.now() - timedelta(days=days)
        deleted = 0

        for file in self.history_dir.glob("*.jsonl"):
            mtime = datetime.fromtimestamp(file.stat().st_mtime)
            if mtime < cutoff:
                file.unlink()
                deleted += 1
                logger.info(f"Deleted old history file: {file.name}")

        return deleted

    # ==================== 容量限制清理 ====================

    # 配置常量
    MAX_HISTORY_DAYS = 30  # 最多保留 30 天
    MAX_HISTORY_FILES = 1000  # 最多保留 1000 个文件
    MAX_HISTORY_SIZE_MB = 500  # 最多占用 500MB

    def cleanup_history(self) -> dict:
        """
        清理历史对话，防止磁盘爆炸

        策略（按优先级）:
        1. 删除超过 MAX_HISTORY_DAYS 天的文件
        2. 如果文件数超过 MAX_HISTORY_FILES，删除最旧的
        3. 如果总大小超过 MAX_HISTORY_SIZE_MB，删除最旧的

        Returns:
            清理统计 {"by_age": n, "by_count": n, "by_size": n}
        """
        deleted = {"by_age": 0, "by_count": 0, "by_size": 0}

        # 1. 按天数清理
        deleted["by_age"] = self.cleanup_old_history(days=self.MAX_HISTORY_DAYS)

        # 获取所有历史文件，按修改时间排序（最旧的在前）
        files = sorted(self.history_dir.glob("*.jsonl"), key=lambda f: f.stat().st_mtime)

        # 2. 按文件数清理
        if len(files) > self.MAX_HISTORY_FILES:
            to_delete = files[: len(files) - self.MAX_HISTORY_FILES]
            for f in to_delete:
                try:
                    f.unlink()
                    deleted["by_count"] += 1
                    logger.debug(f"Deleted history file (by count): {f.name}")
                except Exception as e:
                    logger.error(f"Failed to delete {f.name}: {e}")

            # 更新文件列表
            files = files[len(to_delete) :]

        # 3. 按大小清理
        max_size = self.MAX_HISTORY_SIZE_MB * 1024 * 1024
        total_size = sum(f.stat().st_size for f in files)

        while total_size > max_size and files:
            f = files.pop(0)
            try:
                file_size = f.stat().st_size
                f.unlink()
                total_size -= file_size
                deleted["by_size"] += 1
                logger.debug(f"Deleted history file (by size): {f.name}")
            except Exception as e:
                logger.error(f"Failed to delete {f.name}: {e}")

        total_deleted = sum(deleted.values())
        if total_deleted > 0:
            logger.info(f"History cleanup completed: {deleted}")

        return deleted

    def get_history_stats(self) -> dict:
        """
        获取历史对话统计信息

        Returns:
            统计信息字典
        """
        files = list(self.history_dir.glob("*.jsonl"))
        total_size = sum(f.stat().st_size for f in files)

        return {
            "file_count": len(files),
            "total_size_mb": round(total_size / 1024 / 1024, 2),
            "max_files": self.MAX_HISTORY_FILES,
            "max_size_mb": self.MAX_HISTORY_SIZE_MB,
            "max_days": self.MAX_HISTORY_DAYS,
        }
