import uuid
import json
import asyncio
from astrbot.core import sp
from typing import Dict, List
from astrbot.core.db import BaseDatabase
from astrbot.core.db.po import Conversation


class ConversationManager:
    """负责管理会话与 LLM 的对话，某个会话当前正在用哪个对话。"""

    def __init__(self, db_helper: BaseDatabase):
        self.session_conversations: Dict[str, str] = sp.get("session_conversation", {})
        self.db = db_helper
        self.save_interval = 60  # 每 60 秒保存一次
        self._start_periodic_save()

    def _start_periodic_save(self):
        asyncio.create_task(self._periodic_save())

    async def _periodic_save(self):
        while True:
            await asyncio.sleep(self.save_interval)
            self._save_to_storage()

    def _save_to_storage(self):
        sp.put("session_conversation", self.session_conversations)

    async def new_conversation(self, unified_msg_origin: str) -> str:
        """新建对话，并将当前会话的对话转移到新对话"""
        conversation_id = str(uuid.uuid4())
        self.db.new_conversation(user_id=unified_msg_origin, cid=conversation_id)
        self.session_conversations[unified_msg_origin] = conversation_id
        sp.put("session_conversation", self.session_conversations)
        return conversation_id

    async def switch_conversation(self, unified_msg_origin: str, conversation_id: str):
        """切换会话的对话"""
        self.session_conversations[unified_msg_origin] = conversation_id
        sp.put("session_conversation", self.session_conversations)

    async def delete_conversation(
        self, unified_msg_origin: str, conversation_id: str = None
    ):
        """删除会话的对话，当 conversation_id 为 None 时删除会话当前的对话"""
        conversation_id = self.session_conversations.get(unified_msg_origin)
        if conversation_id:
            self.db.delete_conversation(user_id=unified_msg_origin, cid=conversation_id)
            del self.session_conversations[unified_msg_origin]
            sp.put("session_conversation", self.session_conversations)

    async def get_curr_conversation_id(self, unified_msg_origin: str) -> str:
        """获取会话当前的对话 ID"""
        return self.session_conversations.get(unified_msg_origin, None)

    async def get_conversation(
        self, unified_msg_origin: str, conversation_id: str
    ) -> Conversation:
        """获取会话的对话"""
        return self.db.get_conversation_by_user_id(unified_msg_origin, conversation_id)

    async def get_conversations(self, unified_msg_origin: str) -> List[Conversation]:
        """获取会话的所有对话"""
        return self.db.get_conversations(unified_msg_origin)

    async def update_conversation(
        self, unified_msg_origin: str, conversation_id: str, history: List[Dict]
    ):
        """更新会话的对话"""
        if conversation_id:
            self.db.update_conversation(
                user_id=unified_msg_origin,
                cid=conversation_id,
                history=json.dumps(history),
            )

    async def update_conversation_title(self, unified_msg_origin: str, title: str):
        """更新会话的对话标题"""
        conversation_id = self.session_conversations.get(unified_msg_origin)
        if conversation_id:
            self.db.update_conversation_title(
                user_id=unified_msg_origin, cid=conversation_id, title=title
            )

    async def update_conversation_persona_id(
        self, unified_msg_origin: str, persona_id: str
    ):
        """更新会话的对话 Persona ID"""
        conversation_id = self.session_conversations.get(unified_msg_origin)
        if conversation_id:
            self.db.update_conversation_persona_id(
                user_id=unified_msg_origin, cid=conversation_id, persona_id=persona_id
            )

    async def get_human_readable_context(
        self, unified_msg_origin, conversation_id, page=1, page_size=10
    ):
        conversation = await self.get_conversation(unified_msg_origin, conversation_id)
        history = json.loads(conversation.history)

        contexts = []
        temp_contexts = []
        for record in history:
            if record["role"] == "user":
                temp_contexts.append(f"User: {record['content']}")
            elif record["role"] == "assistant":
                temp_contexts.append(f"Assistant: {record['content']}")
                contexts.insert(0, temp_contexts)
                temp_contexts = []

        # 展平 contexts 列表
        contexts = [item for sublist in contexts for item in sublist]

        # 计算分页
        paged_contexts = contexts[(page - 1) * page_size : page * page_size]
        total_pages = len(contexts) // page_size
        if len(contexts) % page_size != 0:
            total_pages += 1

        return paged_contexts, total_pages