"""对话记忆管理模块

管理 Agent 的对话历史：
- 短期记忆（当前会话）
- 消息历史管理
- 对话摘要（使用 LLM）
- 持久化存储（SQLite）
"""

import os
import sqlite3
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional

from langchain_anthropic import ChatAnthropic

from config.agent import config
from .prompts import get_loader


def _generate_session_id() -> str:
    """生成UUID会话ID"""
    return str(uuid.uuid4())


# 公开版本
def generate_session_id() -> str:
    """生成UUID会话ID"""
    return _generate_session_id()


# 默认数据库路径
DEFAULT_DB_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "db", "sqlite", "memory.db")


@dataclass
class Message:
    """单条消息"""

    role: str  # "user" | "assistant" | "system" | "tool"
    content: str
    timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
    tool_name: Optional[str] = None  # 如果是工具调用结果
    tool_result: Optional[str] = None


class ConversationMemory:
    """对话记忆管理

    支持：
    - 消息历史管理
    - LLM 摘要生成
    - SQLite 持久化存储
    """

    def __init__(
        self,
        max_messages: int = None,
        summary_threshold: int = None,
        session_id: str = "default",
        db_path: str = None,
        llm: ChatAnthropic = None,
    ):
        """
        初始化对话记忆

        Args:
            max_messages: 最多保留的消息数
            summary_threshold: 超过此数量时进行摘要
            session_id: 会话 ID（用于区分不同会话的记忆）
            db_path: SQLite 数据库路径
            llm: 用于摘要的语言模型
        """
        self.max_messages = max_messages or config.memory_max_messages
        self.summary_threshold = summary_threshold or config.memory_summary_threshold
        self.session_id = session_id
        self.db_path = db_path or DEFAULT_DB_PATH
        self._llm = llm
        self.messages: List[Message] = []
        self.summary: Optional[str] = None

        # 初始化数据库
        self._init_db()

    @property
    def llm(self) -> ChatAnthropic:
        """获取 LLM"""
        if self._llm is None:
            self._llm = ChatAnthropic(
                model=config.agent_model,
                temperature=0.3,  # 摘要使用较低温度
                anthropic_api_key=config.anthropic_api_key,
                base_url=config.anthropic_base_url,
            )
        return self._llm

    def _init_db(self):
        """初始化数据库表"""
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        with sqlite3.connect(self.db_path) as conn:
            conn.execute("""
                CREATE TABLE IF NOT EXISTS sessions (
                    id INTEGER PRIMARY KEY,
                    session_id TEXT UNIQUE NOT NULL,
                    summary TEXT,
                    source TEXT DEFAULT 'cli',  -- 'cli', 'feishu', 'api'
                    directory TEXT,  -- CLI 启动目录
                    user_id TEXT,  -- 飞书用户ID
                    chat_id TEXT,  -- 飞书群ID
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            conn.execute("""
                CREATE TABLE IF NOT EXISTS messages (
                    id INTEGER PRIMARY KEY,
                    session_id TEXT NOT NULL,
                    role TEXT NOT NULL,
                    content TEXT NOT NULL,
                    timestamp TEXT NOT NULL,
                    tool_name TEXT,
                    tool_result TEXT,
                    FOREIGN KEY (session_id) REFERENCES sessions(session_id)
                )
            """)

            # 迁移：为已有数据库添加新字段
            try:
                conn.execute("ALTER TABLE sessions ADD COLUMN source TEXT DEFAULT 'cli'")
            except sqlite3.OperationalError:
                pass  # 字段已存在

            try:
                conn.execute("ALTER TABLE sessions ADD COLUMN directory TEXT")
            except sqlite3.OperationalError:
                pass

            try:
                conn.execute("ALTER TABLE sessions ADD COLUMN user_id TEXT")
            except sqlite3.OperationalError:
                pass

            try:
                conn.execute("ALTER TABLE sessions ADD COLUMN chat_id TEXT")
            except sqlite3.OperationalError:
                pass

            conn.commit()

    def _get_connection(self):
        """获取数据库连接"""
        return sqlite3.connect(self.db_path)

    def add_user_message(self, content: str):
        """添加用户消息"""
        self.messages.append(Message(role="user", content=content))

    def add_assistant_message(self, content: str):
        """添加助手消息"""
        self.messages.append(Message(role="assistant", content=content))

    def add_tool_result(self, tool_name: str, result: str):
        """添加工具调用结果"""
        self.messages.append(
            Message(role="tool", content=result, tool_name=tool_name, tool_result=result)
        )

    def add_system_message(self, content: str):
        """添加系统消息"""
        self.messages.append(Message(role="system", content=content))

    def get_recent_messages(self, n: int = None) -> List[Message]:
        """获取最近 n 条消息"""
        n = n or self.max_messages
        return self.messages[-n:] if len(self.messages) > n else self.messages

    def get_messages_for_context(self) -> List[Dict[str, Any]]:
        """
        获取用于构建上下文的消息列表

        格式符合 OpenAI/Claude 消息格式
        """
        messages = []

        # 添加摘要（如果有）
        if self.summary:
            messages.append({
                "role": "system",
                "content": f"之前对话摘要: {self.summary}"
            })

        # 添加最近的消息
        for msg in self.get_recent_messages():
            if msg.role == "tool":
                messages.append({
                    "role": "user",  # 工具结果作为 user 消息
                    "content": f"[{msg.tool_name}]: {msg.tool_result}"
                })
            else:
                messages.append({
                    "role": msg.role,
                    "content": msg.content
                })

        return messages

    def should_summarize(self) -> bool:
        """是否应该进行摘要"""
        return len(self.messages) > self.summary_threshold

    def summarize(self) -> str:
        """
        使用 LLM 生成对话摘要

        Returns:
            生成的摘要文本
        """
        if not self.messages:
            return ""

        # 构建对话历史文本
        conversation_text = "\n".join([
            f"{msg.role}: {msg.content}"
            for msg in self.messages
        ])

        # 摘要提示词
        loader = get_loader()
        default_prompt = f"""请简洁地总结以下对话的要点，保留关键信息和用户意图：

{conversation_text}

摘要："""
        prompt = loader.get("agent_summary", default_prompt, conversation=conversation_text)

        try:
            from langchain_core.messages import HumanMessage
            response = self.llm.invoke([HumanMessage(content=prompt)])

            # 处理 content 可能是字符串或列表的情况
            content = response.content
            if isinstance(content, list):
                # 提取所有 text 类型的块
                self.summary = " ".join([
                    block.get("text", "")
                    for block in content
                    if isinstance(block, dict) and block.get("type") == "text"
                ]).strip()
            else:
                self.summary = content.strip()

            return self.summary
        except Exception as e:
            import logging
            logging.getLogger("memory").error(f"摘要生成失败: {e}")
            return ""

    def save(self, source: str = None, directory: str = None, user_id: str = None, chat_id: str = None):
        """保存当前会话到数据库

        Args:
            source: 会话来源 ('cli', 'feishu', 'api')
            directory: CLI 启动目录
            user_id: 飞书用户ID
            chat_id: 飞书群ID
        """
        with self._get_connection() as conn:
            # 保存或更新会话
            conn.execute("""
                INSERT INTO sessions (session_id, summary, source, directory, user_id, chat_id, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
                ON CONFLICT(session_id) DO UPDATE SET
                    summary=excluded.summary,
                    source=COALESCE(excluded.source, source),
                    directory=COALESCE(excluded.directory, directory),
                    user_id=COALESCE(excluded.user_id, user_id),
                    chat_id=COALESCE(excluded.chat_id, chat_id),
                    updated_at=excluded.updated_at
            """, (
                self.session_id,
                self.summary,
                source or 'cli',
                directory,
                user_id,
                chat_id,
                datetime.now().isoformat()
            ))

            # 删除旧消息
            conn.execute(
                "DELETE FROM messages WHERE session_id = ?",
                (self.session_id,)
            )

            # 保存新消息
            for msg in self.messages:
                conn.execute("""
                    INSERT INTO messages (session_id, role, content, timestamp, tool_name, tool_result)
                    VALUES (?, ?, ?, ?, ?, ?)
                """, (
                    self.session_id,
                    msg.role,
                    msg.content,
                    msg.timestamp,
                    msg.tool_name,
                    msg.tool_result,
                ))

            conn.commit()

    def load(self, session_id: str = None):
        """
        从数据库加载会话记忆

        Args:
            session_id: 要加载的会话 ID（默认当前会话）
        """
        session_id = session_id or self.session_id
        self.session_id = session_id

        with self._get_connection() as conn:
            # 加载会话摘要
            cursor = conn.execute(
                "SELECT summary FROM sessions WHERE session_id = ?",
                (session_id,)
            )
            row = cursor.fetchone()
            if row:
                self.summary = row[0]
            else:
                self.summary = None

            # 加载消息
            self.messages.clear()
            cursor = conn.execute("""
                SELECT role, content, timestamp, tool_name, tool_result
                FROM messages WHERE session_id = ?
                ORDER BY timestamp ASC
            """, (session_id,))

            for row in cursor:
                self.messages.append(Message(
                    role=row[0],
                    content=row[1],
                    timestamp=row[2],
                    tool_name=row[3],
                    tool_result=row[4],
                ))

    def clear(self):
        """清空内存中的记忆"""
        self.messages.clear()
        self.summary = None

    def clear_all(self):
        """清空记忆（内存 + 数据库）"""
        self.clear()
        with self._get_connection() as conn:
            conn.execute("DELETE FROM messages WHERE session_id = ?", (self.session_id,))
            conn.execute("DELETE FROM sessions WHERE session_id = ?", (self.session_id,))
            conn.commit()

    def get_stats(self) -> Dict[str, Any]:
        """获取记忆统计信息"""
        return {
            "session_id": self.session_id,
            "message_count": len(self.messages),
            "has_summary": self.summary is not None,
            "summary_length": len(self.summary) if self.summary else 0,
        }

    def __len__(self) -> int:
        return len(self.messages)

    def __bool__(self) -> bool:
        """对象始终存在，即使没有消息"""
        return True

    @staticmethod
    def find_session_by_directory(directory: str, db_path: str = None) -> Optional[str]:
        """根据目录查找最近的会话ID

        Args:
            directory: CLI 启动目录
            db_path: 数据库路径

        Returns:
            会话ID，如果不存在则返回 None
        """
        db_path = db_path or DEFAULT_DB_PATH
        if not os.path.exists(db_path):
            return None

        with sqlite3.connect(db_path) as conn:
            cursor = conn.execute("""
                SELECT session_id FROM sessions
                WHERE source = 'cli' AND directory = ?
                ORDER BY updated_at DESC
                LIMIT 1
            """, (directory,))
            row = cursor.fetchone()
            return row[0] if row else None

    @staticmethod
    def find_session_by_user_id(user_id: str, db_path: str = None) -> Optional[str]:
        """根据飞书用户ID查找最近的会话ID

        Args:
            user_id: 飞书用户 open_id
            db_path: 数据库路径

        Returns:
            会话ID，如果不存在则返回 None
        """
        db_path = db_path or DEFAULT_DB_PATH
        if not os.path.exists(db_path):
            return None

        with sqlite3.connect(db_path) as conn:
            cursor = conn.execute("""
                SELECT session_id FROM sessions
                WHERE source = 'feishu' AND user_id = ?
                ORDER BY updated_at DESC
                LIMIT 1
            """, (user_id,))
            row = cursor.fetchone()
            return row[0] if row else None

    @staticmethod
    def find_session_by_chat_id(chat_id: str, db_path: str = None) -> Optional[str]:
        """根据飞书群ID查找最近的会话ID

        Args:
            chat_id: 飞书群 chat_id
            db_path: 数据库路径

        Returns:
            会话ID，如果不存在则返回 None
        """
        db_path = db_path or DEFAULT_DB_PATH
        if not os.path.exists(db_path):
            return None

        with sqlite3.connect(db_path) as conn:
            cursor = conn.execute("""
                SELECT session_id FROM sessions
                WHERE source = 'feishu' AND chat_id = ?
                ORDER BY updated_at DESC
                LIMIT 1
            """, (chat_id,))
            row = cursor.fetchone()
            return row[0] if row else None

    @classmethod
    def get_or_create_session(
        cls,
        source: str,
        db_path: str = None,
        directory: str = None,
        user_id: str = None,
        chat_id: str = None,
    ) -> tuple:
        """根据标识符查找或创建会话

        Args:
            source: 会话来源 ('cli', 'feishu', 'api')
            db_path: 数据库路径
            directory: CLI 启动目录
            user_id: 飞书用户ID
            chat_id: 飞书群ID

        Returns:
            (session_id, is_new) 元组
        """
        db_path = db_path or DEFAULT_DB_PATH

        # 1. 根据 source 和标识符查找现有会话
        session_id = None
        if source == 'cli' and directory:
            session_id = cls.find_session_by_directory(directory, db_path)
        elif source == 'feishu':
            if user_id:
                session_id = cls.find_session_by_user_id(user_id, db_path)
            elif chat_id:
                session_id = cls.find_session_by_chat_id(chat_id, db_path)

        # 2. 不存在则创建新会话
        is_new = False
        if not session_id:
            session_id = _generate_session_id()
            is_new = True

            # 保存新会话到数据库
            with sqlite3.connect(db_path) as conn:
                conn.execute("""
                    INSERT INTO sessions (session_id, source, directory, user_id, chat_id, updated_at)
                    VALUES (?, ?, ?, ?, ?, datetime('now'))
                """, (
                    session_id,
                    source,
                    directory,
                    user_id if source == 'feishu' else None,
                    chat_id if source == 'feishu' else None,
                ))
                conn.commit()

        return session_id, is_new

    def __repr__(self) -> str:
        return f"ConversationMemory(messages={len(self.messages)}, summary={self.summary is not None})"
