from __future__ import annotations

import asyncio
import contextlib
import copy
from asyncio import Task
from collections import defaultdict
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from functools import wraps
from io import StringIO
from types import TracebackType
from typing import TYPE_CHECKING, Any, TypeVar
from uuid import uuid4

import aiologic
from jinja2 import Template
from pydantic import BaseModel, Field
from pytz import utc
from typing_extensions import Self

from amrita_core.agent.context import StrategyContext
from amrita_core.agent.strategy import AgentStrategy, NoExceptionHandler
from amrita_core.builtins.agent import ReActAgentStrategy
from amrita_core.consts import ABSTRACT_INSTRUCTION, DEFAULT_TEMPLATE
from amrita_core.hook.exception import FallbackFailed
from amrita_core.preset import PresetManager
from amrita_core.sessions import SessionData
from amrita_core.streaming import SuspendObjectStream
from amrita_core.utils import gather_usage, get_current_datetime_timestamp

from .config import AmritaConfig, get_config
from .hook.event import CompletionEvent, FallbackContext, PreCompletionEvent
from .hook.matcher import MatcherManager
from .libchat import (
    RESPONSE_TYPE,
    call_completion,
    get_last_response,
    get_tokens,
    text_generator,
)
from .logging import debug_log, logger
from .protocol import MessageContent, MessageWithMetadata
from .sessions import SessionsManager
from .tokenizer import hybrid_token_count
from .types import (
    CONTENT_LIST_TYPE,
    CT_MAP,
    USER_INPUT,
    Content,
    ImageContent,
    Message,
    ModelPreset,
    SendMessageWrap,
    TextContent,
    ToolResult,
    UniResponse,
    UniResponseUsage,
)
from .types import (
    MemoryModel as Memory,
)

if TYPE_CHECKING:
    from .sessions import SessionData

# Global lock for thread-safe operations in the chat manager
LOCK = aiologic.Lock()


RESPONSE_CALLBACK_TYPE = Callable[[RESPONSE_TYPE], Awaitable[Any]] | None

# Type vars
FUNC_RET_T = TypeVar("FUNC_RET_T")


# Built-in break points
class SuspendEnum(str, Enum):
    MEMORY = "ChatObject::memory_limiting"
    SINGLE_TOOL = "ChatObject::single_tool_call"
    PRECOMPLE = "matcher_call::pre_completion"
    COMPLE = "matcher_call::post_completion"


class ChatObjectMeta(BaseModel):
    """Metadata model for chat object

    Used to store identification, events, and time information for the chat object.
    """

    stream_id: str  # Chat stream ID
    session_id: str  # Session ID
    user_input: list[TextContent | ImageContent] | str
    time: datetime = Field(default_factory=datetime.now)  # Creation time
    last_call: datetime = Field(default_factory=datetime.now)  # Last call time


class MemoryLimiter:
    """Context processor

    This class is responsible for handling context memory length and token count limits,
    ensuring the chat context remains within predefined constraints by summarizing
    context and removing messages to avoid exceeding the model processing capacity.
    """

    config: AmritaConfig  # Configuration object
    usage: UniResponseUsage | None = None  # Token usage, initially None
    _train: Message[str]  # Training data (system prompts)
    _dropped_messages: list[Message[str] | ToolResult]  # List of removed messages
    _copied_messages: Memory  # Original message copies (for rollback on exceptions)
    _abstract_instruction = ABSTRACT_INSTRUCTION

    def __init__(
        self,
        memory: Memory,
        train: dict[str, str] | Message[str],
        config: AmritaConfig | None = None,
        abstract_instruction: str | None = None,
    ) -> None:
        """Initialize context processor

        Args:
            memory: Memory model to process
            train: Training data (system prompts)
        """
        self.memory: Memory = memory
        self.config = config or get_config()
        self._train = (
            train if isinstance(train, Message) else Message[str].model_validate(train)
        )
        self._abstract_instruction = abstract_instruction or self._abstract_instruction

    async def __aenter__(self) -> Self:
        """Async context manager entry, initialize processing state

        Returns:
            Return instance for use
        """
        self._dropped_messages = []
        self._copied_messages = copy.deepcopy(self.memory)
        logger.debug(
            f"MemoryLimiter initialized, message count: {len(self.memory.messages)}"
        )
        return self

    @classmethod
    def set_abstract_instruction(cls, instruction: str):
        """Set abstract instruction

        Args:
            instruction: Abstract instruction
        """
        if not isinstance(instruction, str):
            raise TypeError("Instruction must be a string")
        elif not instruction:
            raise ValueError("Instruction cannot be empty")
        cls._abstract_instruction = instruction

    @classmethod
    def get_abstract_instruction(cls) -> str:
        """Get abstract instruction

        Returns:
            Abstract instruction
        """
        return cls._abstract_instruction

    @classmethod
    def reset_abstract_instruction(cls):
        """Reset abstract instruction"""
        cls._abstract_instruction = ABSTRACT_INSTRUCTION

    async def _make_abstract(self):
        """Generate context summary

        By calling LLM to summarize all message content in the current memory into a brief content,
        to reduce context length while preserving key information.
        """
        logger.debug("Starting context summarization..")
        proportion = self.config.llm.memory_abstract_proportion  # Summary proportion
        dropped_part: CONTENT_LIST_TYPE = copy.deepcopy(self._dropped_messages)
        index = int(len(self.memory.messages) * proportion) - len(dropped_part)
        if index < 0:
            index = 0
        if index:
            idx = 0
            while idx < len(self.memory.messages):
                element = self.memory.messages[idx]
                dropped_part.append(element)
                if getattr(element, "tool_calls", None) is not None:
                    # This is an assistant message that initiated tool calls
                    # Include all subsequent consecutive tool messages
                    next_idx = idx + 1
                    while next_idx < len(self.memory.messages):
                        next_element = self.memory.messages[next_idx]
                        # Check if this is a tool message (role == 'tool')
                        if getattr(next_element, "role", None) == "tool":
                            dropped_part.append(next_element)
                            next_idx += 1
                        else:
                            break
                    # Update idx to the position after the last processed tool message
                    idx = next_idx
                else:
                    idx += 1
                if idx >= index:
                    break
            self.memory.messages = self.memory.messages[idx:]
        if dropped_part:
            msg_list: CONTENT_LIST_TYPE = [
                Message[str](role="system", content=self._abstract_instruction),
                Message[str](
                    role="user",
                    content=(
                        "Message list:\n```text\n".join(
                            [
                                f"{it}\n"
                                for it in text_generator(
                                    dropped_part,
                                    split_role=True,
                                )
                            ]
                        )
                        + "\n```"
                    ),
                ),
            ]
            logger.debug("Performing context summarization...")
            response = await get_last_response(call_completion(msg_list))
            usage = get_tokens(
                msg_list, response
            )  # Well, this is just a rough calculation.
            self.usage = usage
            logger.debug(f"Context summary received: {response.content}")
            self.memory.abstract = response.content
            logger.debug("Context summarization completed")
        else:
            logger.debug("Context summarization skipped")

    def _drop_message(self):
        """Remove the oldest message from memory and add it to dropped messages list.

        This method removes the first message from the memory and adds it to the
        dropped messages list. If the next message is a tool message, it is also
        removed and added to the dropped messages list.
        """
        data = self.memory
        if len(data.messages) <= 1:
            return
        self._dropped_messages.append(data.messages.pop(0))
        if data.messages[0].role == "tool":
            while data.messages and data.messages[0].role == "tool":
                self._dropped_messages.append(data.messages.pop(0))

    async def run_enforce(self):
        """Execute memory limitation processing

        Execute memory length limitation and token count limitation in sequence,
        ensuring the chat context stays within predefined ranges.
        This method must be used within an async context manager.

        Raises:
            RuntimeError: Thrown when not used in an async context manager
        """
        logger.debug("Starting memory limitation processing..")
        if not hasattr(self, "_dropped_messages") and not hasattr(
            self, "_copied_messages"
        ):
            raise RuntimeError(
                "MemoryLimiter is not initialized, please use `async with MemoryLimiter(memory)` before calling."
            )
        await self._limit_length()
        await self._limit_tokens()
        if self.config.llm.enable_memory_abstract and self._dropped_messages:
            await self._make_abstract()
        logger.debug("Memory limitation processing completed")

    async def _limit_length(self):
        """Control memory length, remove old messages that exceed the limit, remove unsupported messages."""
        logger.debug("Starting memory length limitation..")
        is_multimodal = self.config.llm.enable_multi_modal
        data: Memory = self.memory

        # Process multimodal messages when needed
        for message in data.messages:
            if (
                isinstance(message.content, list)
                and not is_multimodal
                and message.role == "user"
            ):
                message_text = ""
                for content_part in message.content:
                    if isinstance(content_part, dict):
                        validator = CT_MAP.get(content_part["type"])
                        if not validator:
                            raise ValueError(
                                f"Invalid content type: {content_part['type']}"
                            )
                        content_part: Content = validator.model_validate(content_part)
                    if content_part["type"] == "text":
                        message_text += content_part["text"]
                message.content = message_text

        # Enforce memory length limit
        initial_count = len(data.messages)
        while len(data.messages) > 1:
            if data.messages[0].role == "tool":
                data.messages.pop(0)
            elif len(data.messages) > self.config.llm.memory_length_limit:
                self._drop_message()
            else:
                break
        final_count = len(data.messages)
        logger.debug(
            f"Memory length limitation completed, removed {initial_count - final_count} messages"
        )

    async def _limit_tokens(self):
        """Control token count, remove old messages that exceed the limit

        Calculate the token count of the current message list, when exceeding the configured session max token limit,
        gradually remove the earliest messages until satisfying the token count limit.
        """

        def get_token(memory: CONTENT_LIST_TYPE) -> int:
            """Calculate the total token count for a given message list

            Args:
                memory: List of messages to calculate token count for

            Returns:
                Total token count for the messages
            """
            tk_tmp: int = 0
            for msg in text_generator(memory, full_message=True):
                tk_tmp += hybrid_token_count(
                    msg,
                    self.config.llm.tokens_count_mode,
                )
            return tk_tmp

        train = self._train
        train_model = Message.model_validate(train)
        data = self.memory
        logger.debug("Starting token count limitation..")
        memory_l: CONTENT_LIST_TYPE = [train_model, *data.messages]
        if not self.config.llm.enable_tokens_limit:
            logger.debug("Token limitation disabled, skipping processing")
            return
        prompt_length = hybrid_token_count(train.content)
        if prompt_length > self.config.llm.session_tokens_windows:
            logger.warning(
                f"Prompt size too large! It's {prompt_length}>{self.config.llm.session_tokens_windows}! Please adjusts the prompt or settings!"
            )
            return
        tk_tmp: int = get_token(memory_l)

        initial_count = len(data.messages)
        while tk_tmp > self.config.llm.session_tokens_windows:
            if len(data.messages) > 1:
                self._drop_message()
            else:
                break

            tk_tmp: int = get_token(memory_l)
            memory_l = [train_model, *data.messages]
            await asyncio.sleep(
                0
            )  # CPU intensive tasks may cause performance issues, yielding control here
        final_count = len(data.messages)
        logger.debug(
            f"Token count limitation completed, removed {initial_count - final_count} messages"
        )
        logger.debug(f"Final token count: {tk_tmp}")

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Async context manager exit, handle rollback in case of exceptions

        In case of exceptions, restore messages to the state before processing,
        ensuring data consistency.

        Args:
            exc_type: Exception type
            exc_val: Exception value
            exc_tb: Exception traceback
        """
        del exc_val, exc_tb
        if exc_type is not None:
            logger.warning("An exception occurred, rolling back messages...")
            self.memory.messages = self._copied_messages.messages
            return


class ChatObject(SuspendObjectStream[RESPONSE_TYPE]):
    """Chat processing object - The minimal unit of chat processing.

    This class is responsible for processing a single chat session, including message receiving,
    context management, model calling, and response sending.
    """

    stream_id: str  # Chat object ID
    timestamp: str  # Timestamp (for LLM)
    time: datetime  # Time
    end_at: datetime | None = None
    data: Memory  # (lateinit) Memory file
    user_input: USER_INPUT
    user_message: Message[USER_INPUT]  # (lateinit) User message
    context_wrap: SendMessageWrap  # (lateinit) Context message
    train: Message[str]  # System message
    last_call: datetime  # Last internal function call time
    now_calling: str | None = None  # currently calling function name
    session_id: str  # Session ID
    response: UniResponse[str, None]  # (lateinit) Response
    extra_usage: UniResponseUsage[int]
    preset: ModelPreset  # preset used in this call
    config: AmritaConfig  # config used in this call
    session: SessionData | None  # (lateinit) Session data
    strategy: type[AgentStrategy]
    template: Template
    jinja2_vars: dict[str, Any]  # Vars will be passed to template system.
    _q_tout: float | None
    _is_running: bool = False  # Whether it is running
    _is_done: bool = False  # Whether it has completed
    _task: Task[None]
    _err: BaseException | None = None
    _hook_kwargs: dict[str, Any]
    _hook_args: tuple[Any, ...]

    _raised_exc: tuple[type[BaseException], ...]

    def __init__(
        self,
        train: dict[str, str] | Message[str],
        user_input: USER_INPUT,
        context: Memory | None,
        session_id: str,
        callback: RESPONSE_CALLBACK_TYPE = None,
        config: AmritaConfig | None = None,
        preset: ModelPreset | None = None,
        auto_create_session: bool = False,
        *,
        train_template: Template = DEFAULT_TEMPLATE,
        jinja2_vars: dict[str, Any] | None = None,
        agent_strategy: type[AgentStrategy] = ReActAgentStrategy,
        hook_args: tuple[Any, ...] = (),
        hook_kwargs: dict[str, Any] | None = None,
        exception_ignored: tuple[type[BaseException], ...] = (),
        queue_size: int = 45,
        queue_timeout: float | None = 10.0,
    ) -> None:
        """Initialize a chat object

        Args:
            train (dict[str, str] | Message[str]): Training data (system prompts)
            user_input (USER_INPUT): Input from the user
            context (Memory | None): Memory context for the session
            session_id (str): Unique identifier for the session
            callback (RESPONSE_CALLBACK_TYPE, optional): Callback function to be called when returning response. Defaults to None.
            config (AmritaConfig | None, optional): Config used for this call. Defaults to None.
            preset (ModelPreset | None, optional): Preset used for this call. Defaults to None.
            auto_create_session (bool, optional): Whether to automatically create a session if it does not exist. Defaults to False.
            jinja2_vars (dict[str, Any] | None, optional): Variables to be passed to the template system. Defaults to None.
            train_template (Template, optional): Jinja2 template used to format system message.
            agent_strategy (type[AgentStrategy], optional):  Agent strategy to be used for execution. Defaults to ReActAgentStrategy.
            hook_args (tuple[Any, ...], optional): Arguments could be passed to the Matcher function. Defaults to ().
            hook_kwargs (dict[str, Any] | None, optional): Keyword arguments could be passed to the Matcher function. Defaults to None.
            exception_ignored (tuple[type[BaseException], ...], optional): These exceptions will be raised again if they are raised in the Matcher function. Defaults to ().
            queue_size (int, optional): Maximum number of message chunks to be stored in the queue. Defaults to 45.
        """
        sm = SessionsManager()
        if auto_create_session and not sm.is_session_registered(session_id):
            sm.init_session(session_id)
        self._raised_exc = exception_ignored
        session: SessionData | None = sm.get_session_data(session_id, None)
        self.session = session
        self.train = (
            train if isinstance(train, Message) else Message[str].model_validate(train)
        )
        self.data = context or sm.get_session_data(session_id).memory
        self.session_id = session_id
        # initialize iostream
        super().__init__(
            queue_size=queue_size, callback=callback, queue_timeout=queue_timeout
        )
        # data
        self.user_input = user_input
        self.user_message = Message(role="user", content=user_input)
        self.timestamp = get_current_datetime_timestamp()
        self.time = datetime.now(utc)
        self.config: AmritaConfig = config or (
            session.config if session else get_config()
        )
        self.strategy = agent_strategy
        self.extra_usage = UniResponseUsage(
            prompt_tokens=0, completion_tokens=0, total_tokens=0
        )
        # other
        self.last_call = datetime.now(utc)
        self.preset = preset or (
            session.presets.get_default_preset()
            if session
            else PresetManager().get_default_preset()
        )
        self.template = train_template
        jinja2_vars = jinja2_vars or {}
        if any(name in jinja2_vars for name in ("train", "self", "memory", "chatobj")):
            raise RuntimeError("Received a reserved keyword, please use another name.")
        self.jinja2_vars = jinja2_vars
        # Hook args
        self._hook_args = hook_args
        self._hook_kwargs = hook_kwargs or {}
        # initialize id
        self.stream_id = uuid4().hex

    @staticmethod
    def monitoring(func: Callable[..., Any]):
        """Decorator for monitoring.This decorator will ONLY BE USED in ChatObject."""

        @wraps(func)
        def inner(*args, **kwargs):
            self: ChatObject = args[0]  # This is Self
            self.last_call = datetime.now(utc)
            pev = self.now_calling
            self.now_calling = func.__name__
            try:
                return func(*args, **kwargs)
            finally:
                self.now_calling = pev

        return inner

    async def full_response(self) -> str:
        """Return full response from the queue as a single string.

        Returns:
            Complete response string combining all chunks in the queue
        """
        builder = StringIO()
        async for item in self.get_response_generator():
            if isinstance(item, str):
                builder.write(item)
            elif isinstance(item, MessageContent):
                builder.write(str(item.get_content()))
        return builder.getvalue()

    def get_exception(self) -> BaseException | None:
        """
        Get exceptions that occurred during task execution

        Returns:
            BaseException | None: Returns exception object if an exception occurred during task execution, otherwise returns None
        """
        return self._err

    def is_running(self) -> bool:
        """
        Check if the task is running

        Returns:
            bool: Returns True if the task is running, otherwise returns False
        """
        return self._is_running

    def is_done(self) -> bool:
        """
        Check if the task has completed

        Returns:
            bool: Returns True if the task has completed, otherwise returns False
        """
        return self._is_done

    def terminate(self) -> None:
        """
        Terminate task execution
        Sets the task status to completed and cancels the internal task
        """
        self._is_done = True
        self._is_running = False
        if hasattr(self, "_task") and not self._task.done():
            self._task.cancel()

    def __await__(self):
        """
        Await for task completion
        """
        if not hasattr(self, "_task"):
            raise RuntimeError("ChatObject not running")
        return self._task.__await__()

    async def __aenter__(self) -> Self:
        if not hasattr(self, "_task"):
            raise RuntimeError("ChatObject not running")
        if self._has_consumer:
            raise RuntimeError("ChatObject already has a consumer")
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: TracebackType | None,
    ):
        del exc_tb  # This is unused
        if exc_type is not None:
            self._err = exc_val
        self.terminate()

    def begin(self) -> Self:
        """Start chat object task"""
        if not hasattr(self, "_task"):
            logger.debug("Starting chat object task...")
            self._task = asyncio.create_task(self._entry())
        return self

    @monitoring
    @SuspendObjectStream.suspend
    async def _entry(self) -> None:
        """Call chat object to process messages"""
        if not self._is_running and not self._is_done:
            self.stream_id = uuid4().hex
            logger.debug(f"Starting chat processing, stream ID:{self.stream_id}")

            try:
                self._is_running = True
                await chat_manager.add_chat_object(self)

                await self._run()
            finally:
                self._is_running = False
                self._is_done = True
                self.end_at = datetime.now(utc)
                chat_manager.running_chat_object_id2map.pop(self.stream_id, None)
                chat_manager.clean_obj(self.session_id, 1000)  # To avoid memory leaks
                logger.debug("Chat event processing completed")

        else:
            raise RuntimeError(
                f"ChatObject of {self.stream_id} is already running or done"
            )

    @SuspendObjectStream.suspend
    async def _run(self):
        """Run chat processing flow

        Execute the main logic of message processing, including getting user information,
        processing message content, managing context length and token limitations,
        and sending responses.
        """
        logger.debug("Starting chat processing flow..")
        data = self.data
        config = self.config

        data.messages.append(self.user_message)

        logger.debug(
            f"Added user message to memory, current message count: {len(data.messages)}"
        )
        # train,memory,chatobj(ChatObject),config will be given to Jinja2
        self.train.content = await asyncio.to_thread(
            self.template.render,
            train=self.train,
            memory=self.data,
            chatobj=self,
            config=config,
            **self.jinja2_vars,
        )
        debug_log(self.train.content)
        logger.debug("Starting applying memory limitations..")
        async with MemoryLimiter(self.data, self.train, config=config) as lim:
            await self._wait_for_continue(SuspendEnum.MEMORY.value)
            await lim.run_enforce()

            if abs_usage := lim.usage:
                self.extra_usage = gather_usage(self.extra_usage, abs_usage)
            self.data = lim.memory
        logger.debug("Memory limitation application completed")

        send_messages = self._prepare_send_messages()
        self.context_wrap = SendMessageWrap.validate_messages(send_messages)
        logger.debug(
            f"Preparing sending messages completed, message count: {len(send_messages)}"
        )
        response: UniResponse[str, None] = await self._process_chat(send_messages)
        self.response = response

        logger.debug("Chat processing completed, preparing to send response")
        await self.set_queue_done()

    @SuspendObjectStream.suspend
    async def _run_strategy(self) -> None:
        """Run workflow of strategy given."""

        match self.strategy.get_category():
            case "agent-mixed" | "agent":
                context = (
                    SendMessageWrap.validate_messages([self.train, self.user_message])
                    if self.config.function_config.use_minimal_context
                    else self.context_wrap.copy()
                )
                ctx = StrategyContext(self.user_input, context, self)
                return await self._run_agent(ctx)

            case "rag":
                context = SendMessageWrap.validate_messages(
                    [
                        self.train,
                        self.user_message,
                    ]
                )
            case "workflow":
                context = self.context_wrap.copy()
            case _:
                raise RuntimeError("Invalid agent strategy")
        ctx = StrategyContext(self.user_input, context, self)
        st = self.strategy(ctx)
        try:
            await st.run()
        except Exception as e:
            if isinstance(e, self._raised_exc):
                raise
            with contextlib.suppress(NoExceptionHandler):
                await st.on_exception(e)
        else:
            await st.on_post_process()
        self.context_wrap.extend(ctx.original_context.end_messages)

    @SuspendObjectStream.suspend
    async def _run_agent(self, ctx: StrategyContext) -> None:
        strategy: AgentStrategy = self.strategy(ctx)
        backup: SendMessageWrap = self.context_wrap.copy()
        try:
            for _ in range(1, self.config.function_config.agent_tool_call_limit + 1):
                await self._wait_for_continue(SuspendEnum.SINGLE_TOOL.value)
                if not (await strategy.single_execute()):
                    break
            else:
                await strategy.on_limited()
            await strategy.on_post_process()
            self.context_wrap.extend(strategy.ctx.original_context.end_messages)

        except Exception as e:
            if isinstance(e, self._raised_exc):
                raise
            logger.warning(
                f"ERROR\n{e!s}\n!Failed to call Tools! Continuing with old data..."
            )
            await self.yield_response(
                MessageWithMetadata(
                    content=f"Agent run failed:{e!s}",
                    metadata={"type": "error", "error": e},
                )
            )
            await strategy.on_exception(e)
            self.context_wrap = backup

    @monitoring
    @SuspendObjectStream.suspend
    async def _process_chat(
        self,
        send_messages: CONTENT_LIST_TYPE,
    ) -> UniResponse[str, None]:
        """Call chat model to generate response and trigger related events.

        Args:
            send_messages: Send message list
            extra_usage: Extra token usage information

        Returns:
            Model response
        """

        data = self.data
        logger.debug(
            f"Starting chat processing, sending message count: {len(send_messages)}"
        )

        logger.debug("Triggering matcher functions..")
        messages = self.context_wrap
        chat_event = PreCompletionEvent(
            chat_object=self,
            user_input=self.user_input,
            original_context=messages,
        )
        await self._wait_for_continue(SuspendEnum.PRECOMPLE.value)
        await MatcherManager.trigger_event(
            chat_event,
            self.config,
            self,
            self.preset,
            *self._hook_args,
            session=self.session,
            exception_ignored=self._raised_exc,
            **self._hook_kwargs,
        )
        await self._run_strategy()
        self.data.messages = chat_event.get_context_messages().unwrap(
            exclude_system=True
        )
        send_messages = chat_event.get_context_messages().unwrap()
        logger.debug("Calling chat model..")
        response: UniResponse[str, None] | None = None
        used_preset: set[str] = set()
        for i in range(1, self.config.llm.max_fallbacks + 1):
            try:
                used_preset.add(self.preset.name)
                async for chunk in call_completion(
                    send_messages, config=self.config, preset=self.preset
                ):
                    if isinstance(chunk, UniResponse):
                        response = chunk
                    elif isinstance(chunk, MessageContent | str):
                        await self.yield_response(chunk)
                break
            except Exception as e:
                logger.warning(
                    f"Because of `{e!s}`, LLM request failed, retrying ({i}/{self.config.llm.max_retries})..."
                )
                ctx = FallbackContext(self.preset, e, self.config, self.context_wrap, i)
                await MatcherManager.trigger_event(ctx, ctx.config, (FallbackFailed,))
                if ctx.preset is self.preset:
                    ctx.fail("No preset fallback available, exiting!")
                self.preset = ctx.preset
        else:
            raise FallbackFailed("Max preset fallbacks retries exceeded.")
        if response is None:
            raise RuntimeError("No final response from chat adapter.")
        self.response = response
        logger.debug("Triggering chat events..")
        chat_event = CompletionEvent(self.user_input, messages, self, response.content)
        await self._wait_for_continue(SuspendEnum.COMPLE.value)
        await MatcherManager.trigger_event(
            chat_event,
            self.config,
            self,
            self.preset,
            *self._hook_args,
            session=self.session,
            exception_ignored=self._raised_exc,
            **self._hook_kwargs,
        )
        response.content = chat_event.model_response
        data.messages.append(
            Message[str](
                content=response.content,
                role="assistant",
            )
        )
        logger.debug(
            f"Added assistant response to memory, current message count: {len(data.messages)}"
        )

        logger.debug("Chat processing completed")
        return response

    @monitoring
    def _prepare_send_messages(
        self,
    ) -> list:
        """Prepare message list to send to the chat model, including system prompt data and context.

        Returns:
            Prepared message list to send
        """
        logger.debug("Preparing messages to send..")
        train: Message[str] = Message[str].model_validate(self.train)
        data = self.data
        messages = [train, *copy.deepcopy(data.messages)]
        logger.debug(f"Messages preparation completed, total {len(messages)} messages")
        return messages

    def get_snapshot(self) -> ChatObjectMeta:
        """Get a snapshot of the chat object

        Returns:
            Chat object metadata
        """
        return ChatObjectMeta.model_validate(self, from_attributes=True)


@dataclass
class ChatManager:
    running_chat_object: defaultdict[str, list[ChatObject]] = field(
        default_factory=lambda: defaultdict(list)
    )
    running_chat_object_id2map: dict[str, ChatObjectMeta] = field(default_factory=dict)

    def clean_obj(self, k: str, maxitems: int = 10):
        """
        Clean up running chat objects under the specified key, keeping only the first 10 objects,
        removing any excess unfinished parts

        Args:
            k (tuple[int, bool]): Key value, composed of instance ID and whether it's group chat
            maxitems (int, optional): Maximum number of objects. Defaults to 10.
        """
        objs = self.running_chat_object[k]
        if len(objs) > maxitems:
            dropped_obj = objs[maxitems:]
            objs = [obj for obj in dropped_obj if not obj.is_done()] + objs[:maxitems]
            dropped_obj = [obj for obj in dropped_obj if obj.is_done()]
            for obj in dropped_obj:
                self.running_chat_object_id2map.pop(obj.stream_id, None)
            self.running_chat_object[k] = objs

    def get_all_objs(self) -> list[ChatObjectMeta]:
        """
        Get all running chat object metadata

        Returns:
            list[ChatObjectMeta]: List of all running chat object metadata
        """
        return list(self.running_chat_object_id2map.values())

    def get_objs(self, session_id: str) -> list[ChatObject]:
        """
        Get the corresponding list of chat objects based on the session ID

        Args:
            session_id (str): User session ID

        Returns:
            list[ChatObject]: List of chat objects
        """
        return self.running_chat_object[session_id]

    async def clean_chat_objects(self, maxitems: int = 10) -> None:
        """
        Asynchronously clean up all running chat objects, limiting the number of objects for each key to no more than 10
        """
        async with LOCK:
            for key in self.running_chat_object.keys():
                self.clean_obj(key, maxitems)

    async def add_chat_object(self, chat_object: ChatObject) -> None:
        """
        Add a new chat object to the running list

        Args:
            chat_object (ChatObject): Chat object instance
        """
        async with LOCK:
            meta: ChatObjectMeta = chat_object.get_snapshot()
            self.running_chat_object_id2map[chat_object.stream_id] = meta
            key = chat_object.session_id
            self.running_chat_object[key].insert(0, chat_object)
            self.clean_obj(key)


chat_manager = ChatManager()
