from typing import Dict, TypeVar, List, Optional
import traceback
from json_repair import repair_json
from rich.traceback import install
from src.common.logger import get_logger
from src.llm_models.utils_model import LLMRequest
from src.config.config import global_config
from src.chat.focus_chat.working_memory.memory_item import MemoryItem
import json  # 添加json模块导入


install(extra_lines=3)
logger = get_logger("working_memory")

T = TypeVar("T")


class MemoryManager:
    def __init__(self, chat_id: str):
        """
        初始化工作记忆

        Args:
            chat_id: 关联的聊天ID，用于标识该工作记忆属于哪个聊天
        """
        # 关联的聊天ID
        self._chat_id = chat_id

        # 记忆项列表
        self._memories: List[MemoryItem] = []

        # ID到记忆项的映射
        self._id_map: Dict[str, MemoryItem] = {}

        self.llm_summarizer = LLMRequest(
            model=global_config.model.focus_working_memory,
            temperature=0.3,
            request_type="focus.processor.working_memory",
        )

    @property
    def chat_id(self) -> str:
        """获取关联的聊天ID"""
        return self._chat_id

    @chat_id.setter
    def chat_id(self, value: str):
        """设置关联的聊天ID"""
        self._chat_id = value

    def push_item(self, memory_item: MemoryItem) -> str:
        """
        推送一个已创建的记忆项到工作记忆中

        Args:
            memory_item: 要存储的记忆项

        Returns:
            记忆项的ID
        """
        # 添加到内存和ID映射
        self._memories.append(memory_item)
        self._id_map[memory_item.id] = memory_item

        return memory_item.id

    def get_by_id(self, memory_id: str) -> Optional[MemoryItem]:
        """
        通过ID获取记忆项

        Args:
            memory_id: 记忆项ID

        Returns:
            找到的记忆项，如果不存在则返回None
        """
        memory_item = self._id_map.get(memory_id)
        if memory_item:
            # 检查记忆强度，如果小于1则删除
            if not memory_item.is_memory_valid():
                print(f"记忆 {memory_id} 强度过低 ({memory_item.memory_strength})，已自动移除")
                self.delete(memory_id)
                return None

        return memory_item

    def get_all_items(self) -> List[MemoryItem]:
        """获取所有记忆项"""
        return list(self._id_map.values())

    def find_items(
        self,
        source: Optional[str] = None,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
        memory_id: Optional[str] = None,
        limit: Optional[int] = None,
        newest_first: bool = False,
        min_strength: float = 0.0,
    ) -> List[MemoryItem]:
        """
        按条件查找记忆项

        Args:
            source: 数据来源
            start_time: 开始时间戳
            end_time: 结束时间戳
            memory_id: 特定记忆项ID
            limit: 返回结果的最大数量
            newest_first: 是否按最新优先排序
            min_strength: 最小记忆强度

        Returns:
            符合条件的记忆项列表
        """
        # 如果提供了特定ID，直接查找
        if memory_id:
            item = self.get_by_id(memory_id)
            return [item] if item else []

        results = []

        # 获取所有项目
        items = self._memories

        # 如果需要最新优先，则反转遍历顺序
        if newest_first:
            items_to_check = list(reversed(items))
        else:
            items_to_check = items

        # 遍历项目
        for item in items_to_check:
            # 检查来源是否匹配
            if source is not None and not item.matches_source(source):
                continue

            # 检查时间范围
            if start_time is not None and item.timestamp < start_time:
                continue
            if end_time is not None and item.timestamp > end_time:
                continue

            # 检查记忆强度
            if min_strength > 0 and item.memory_strength < min_strength:
                continue

            # 所有条件都满足，添加到结果中
            results.append(item)

            # 如果达到限制数量，提前返回
            if limit is not None and len(results) >= limit:
                return results

        return results

    async def summarize_memory_item(self, content: str) -> Dict[str, str]:
        """
        使用LLM总结记忆项

        Args:
            content: 需要总结的内容

        Returns:
            包含brief和summary的字典
        """
        prompt = f"""请对以下内容进行总结，总结成记忆，输出两部分：
1. 记忆内容主题（精简，20字以内）：让用户可以一眼看出记忆内容是什么
2. 记忆内容概括：对内容进行概括，保留重要信息，200字以内

内容：
{content}

请按以下JSON格式输出：
{{
  "brief": "记忆内容主题",
  "summary": "记忆内容概括"
}}
请确保输出是有效的JSON格式，不要添加任何额外的说明或解释。
"""
        default_summary = {
            "brief": "主题未知的记忆",
            "summary": "无法概括的记忆内容",
        }

        try:
            # 调用LLM生成总结
            response, _ = await self.llm_summarizer.generate_response_async(prompt)

            # 使用repair_json解析响应
            try:
                # 使用repair_json修复JSON格式
                fixed_json_string = repair_json(response)

                # 如果repair_json返回的是字符串，需要解析为Python对象
                if isinstance(fixed_json_string, str):
                    try:
                        json_result = json.loads(fixed_json_string)
                    except json.JSONDecodeError as decode_error:
                        logger.error(f"JSON解析错误: {str(decode_error)}")
                        return default_summary
                else:
                    # 如果repair_json直接返回了字典对象，直接使用
                    json_result = fixed_json_string

                # 进行额外的类型检查
                if not isinstance(json_result, dict):
                    logger.error(f"修复后的JSON不是字典类型: {type(json_result)}")
                    return default_summary

                # 确保所有必要字段都存在且类型正确
                if "brief" not in json_result or not isinstance(json_result["brief"], str):
                    json_result["brief"] = "主题未知的记忆"

                if "summary" not in json_result or not isinstance(json_result["summary"], str):
                    json_result["summary"] = "无法概括的记忆内容"

                return json_result

            except Exception as json_error:
                logger.error(f"JSON处理失败: {str(json_error)}，将使用默认摘要")
                return default_summary

        except Exception as e:
            logger.error(f"生成总结时出错: {str(e)}")
            return default_summary

    def decay_memory(self, memory_id: str, decay_factor: float = 0.8) -> bool:
        """
        使单个记忆衰减

        Args:
            memory_id: 记忆ID
            decay_factor: 衰减因子(0-1之间)

        Returns:
            是否成功衰减
        """
        memory_item = self.get_by_id(memory_id)
        if not memory_item:
            return False

        # 计算衰减量（当前强度 * (1-衰减因子)）
        old_strength = memory_item.memory_strength
        decay_amount = old_strength * (1 - decay_factor)

        # 更新强度
        memory_item.memory_strength = decay_amount

        return True

    def delete(self, memory_id: str) -> bool:
        """
        删除指定ID的记忆项

        Args:
            memory_id: 要删除的记忆项ID

        Returns:
            是否成功删除
        """
        if memory_id not in self._id_map:
            return False

        # 获取要删除的项
        self._id_map[memory_id]

        # 从内存中删除
        self._memories = [i for i in self._memories if i.id != memory_id]

        # 从ID映射中删除
        del self._id_map[memory_id]

        return True

    def clear(self) -> None:
        """清除所有记忆"""
        self._memories.clear()
        self._id_map.clear()

    async def merge_memories(
        self, memory_id1: str, memory_id2: str, reason: str, delete_originals: bool = True
    ) -> MemoryItem:
        """
        合并两个记忆项

        Args:
            memory_id1: 第一个记忆项ID
            memory_id2: 第二个记忆项ID
            reason: 合并原因
            delete_originals: 是否删除原始记忆，默认为True

        Returns:
            合并后的记忆项
        """
        # 获取两个记忆项
        memory_item1 = self.get_by_id(memory_id1)
        memory_item2 = self.get_by_id(memory_id2)

        if not memory_item1 or not memory_item2:
            raise ValueError("无法找到指定的记忆项")

        # 构建合并提示
        prompt = f"""
请根据以下原因，将两段记忆内容有机合并成一段新的记忆内容。
合并时保留两段记忆的重要信息，避免重复，确保生成的内容连贯、自然。

合并原因：{reason}

记忆1主题：{memory_item1.brief}
记忆1内容：{memory_item1.summary}

记忆2主题：{memory_item2.brief}
记忆2内容：{memory_item2.summary}

请按以下JSON格式输出合并结果：
{{
    "brief": "合并后的主题（20字以内）",
    "summary": "合并后的内容概括（200字以内）"
}}
请确保输出是有效的JSON格式，不要添加任何额外的说明或解释。
"""

        # 默认合并结果
        default_merged = {
            "brief": f"合并：{memory_item1.brief} + {memory_item2.brief}",
            "summary": f"合并的记忆：{memory_item1.summary}\n{memory_item2.summary}",
        }

        try:
            # 调用LLM合并记忆
            response, _ = await self.llm_summarizer.generate_response_async(prompt)

            # 处理LLM返回的合并结果
            try:
                # 修复JSON格式
                fixed_json_string = repair_json(response)

                # 将修复后的字符串解析为Python对象
                if isinstance(fixed_json_string, str):
                    try:
                        merged_data = json.loads(fixed_json_string)
                    except json.JSONDecodeError as decode_error:
                        logger.error(f"JSON解析错误: {str(decode_error)}")
                        merged_data = default_merged
                else:
                    # 如果repair_json直接返回了字典对象，直接使用
                    merged_data = fixed_json_string

                # 确保是字典类型
                if not isinstance(merged_data, dict):
                    logger.error(f"修复后的JSON不是字典类型: {type(merged_data)}")
                    merged_data = default_merged

                if "brief" not in merged_data or not isinstance(merged_data["brief"], str):
                    merged_data["brief"] = default_merged["brief"]

                if "summary" not in merged_data or not isinstance(merged_data["summary"], str):
                    merged_data["summary"] = default_merged["summary"]

            except Exception as e:
                logger.error(f"合并记忆时处理JSON出错: {str(e)}")
                traceback.print_exc()
                merged_data = default_merged
        except Exception as e:
            logger.error(f"合并记忆调用LLM出错: {str(e)}")
            traceback.print_exc()
            merged_data = default_merged

        # 创建新的记忆项
        # 取两个记忆项中更强的来源
        merged_source = (
            memory_item1.from_source
            if memory_item1.memory_strength >= memory_item2.memory_strength
            else memory_item2.from_source
        )

        # 创建新的记忆项
        merged_memory = MemoryItem(
            summary=merged_data["summary"], from_source=merged_source, brief=merged_data["brief"]
        )

        # 记忆强度取两者最大值
        merged_memory.memory_strength = max(memory_item1.memory_strength, memory_item2.memory_strength)

        # 添加到存储中
        self.push_item(merged_memory)

        # 如果需要，删除原始记忆
        if delete_originals:
            self.delete(memory_id1)
            self.delete(memory_id2)

        return merged_memory

    def delete_earliest_memory(self) -> bool:
        """
        删除最早的记忆项

        Returns:
            是否成功删除
        """
        # 获取所有记忆项
        all_memories = self.get_all_items()

        if not all_memories:
            return False

        # 按时间戳排序，找到最早的记忆项
        earliest_memory = min(all_memories, key=lambda item: item.timestamp)

        # 删除最早的记忆项
        return self.delete(earliest_memory.id)
