import copy
import warnings
from collections import OrderedDict, UserDict, UserList, abc
from functools import wraps
from itertools import chain, repeat
from typing import Any, AsyncGenerator, Callable, Dict, Generator, Iterable, List, Mapping, Optional, Tuple, Union

from lagent.agents.aggregator import DefaultAggregator
from lagent.hooks import Hook, RemovableHandle
from lagent.llms import BaseLLM
from lagent.memory import Memory, MemoryManager
from lagent.prompts.parsers import StrParser
from lagent.prompts.prompt_template import PromptTemplate
from lagent.schema import AgentMessage, ModelStatusCode
from lagent.utils import create_object


class Agent:
    """Agent is the basic unit of the system. It is responsible for
    communicating with the LLM, managing the memory, and handling the
    message aggregation and parsing. It can also be extended with hooks

    Args:
        llm (Union[BaseLLM, Dict]): The language model used by the agent.
        template (Union[PromptTemplate, str]): The template used to format the
            messages.
        memory (Dict): The memory used by the agent.
        output_format (Dict): The output format used by the agent.
        aggregator (Dict): The aggregator used by the agent.
        name (Optional[str]): The name of the agent.
        description (Optional[str]): The description of the agent.
        hooks (Optional[Union[List[Dict], Dict]]): The hooks used by the agent.

    Returns:
        AgentMessage: The response message.
    """

    def __init__(
        self,
        llm: Union[BaseLLM, Dict] = None,
        template: Union[PromptTemplate, str, dict, List[dict]] = None,
        memory: Dict = dict(type=Memory),
        output_format: Optional[Dict] = None,
        aggregator: Dict = dict(type=DefaultAggregator),
        name: Optional[str] = None,
        description: Optional[str] = None,
        hooks: Optional[Union[List[Dict], Dict]] = None,
    ):
        self.name = name or self.__class__.__name__
        self.llm: BaseLLM = create_object(llm)
        self.memory: MemoryManager = MemoryManager(memory) if memory else None
        self.output_format: StrParser = create_object(output_format)
        self.template = template
        self.description = description
        self.aggregator: DefaultAggregator = create_object(aggregator)
        self._hooks: Dict[int, Hook] = OrderedDict()
        if hooks:
            for hook in hooks:
                hook = create_object(hook)
                self.register_hook(hook)

    def update_memory(self, message, session_id=0):
        if self.memory:
            self.memory.add(message, session_id=session_id)

    def __call__(self, *message: AgentMessage, session_id=0, **kwargs) -> AgentMessage:
        # message.receiver = self.name
        message = [AgentMessage(sender='user', content=m) if isinstance(m, str) else copy.deepcopy(m) for m in message]
        for hook in self._hooks.values():
            result = hook.before_agent(self, message, session_id)
            if result:
                message = result
        self.update_memory(message, session_id=session_id)
        response_message = self.forward(*message, session_id=session_id, **kwargs)
        if not isinstance(response_message, AgentMessage):
            response_message = AgentMessage(sender=self.name, content=response_message)
        self.update_memory(response_message, session_id=session_id)
        response_message = copy.deepcopy(response_message)
        for hook in self._hooks.values():
            result = hook.after_agent(self, response_message, session_id)
            if result:
                response_message = result
        return response_message

    def forward(self, *message: AgentMessage, session_id=0, **kwargs) -> Union[AgentMessage, str]:
        formatted_messages = self.aggregator.aggregate(
            self.memory.get(session_id), self.name, self.output_format, self.template
        )
        llm_response = self.llm.chat(formatted_messages, **kwargs)
        if self.output_format:
            formatted_messages = self.output_format.parse_response(llm_response)
            return AgentMessage(sender=self.name, content=llm_response, formatted=formatted_messages)
        return llm_response

    def __setattr__(self, __name: str, __value: Any) -> None:
        if isinstance(__value, Agent):
            _agents = getattr(self, '_agents', OrderedDict())
            _agents[__name] = __value
            super().__setattr__('_agents', _agents)
        super().__setattr__(__name, __value)

    def state_dict(self, session_id=0):
        state_dict, stack = {}, [('', self)]
        while stack:
            prefix, node = stack.pop()
            key = prefix + 'memory'
            if node.memory is not None:
                if session_id not in node.memory.memory_map:
                    warnings.warn(f'No session id {session_id} in {key}')
                memory = node.memory.get(session_id)
                state_dict[key] = memory and memory.save() or []
            if hasattr(node, '_agents'):
                for name, value in reversed(node._agents.items()):
                    stack.append((prefix + name + '.', value))
        return state_dict

    def load_state_dict(self, state_dict: Dict, session_id=0):
        _state_dict = self.state_dict()
        missing_keys = set(_state_dict) - set(state_dict)
        if missing_keys:
            raise KeyError(f'Missing keys: {missing_keys}')
        extra_keys = set(state_dict) - set(_state_dict)
        if extra_keys:
            warnings.warn(f'Mismatch keys which are not used: {extra_keys}')
        for key in _state_dict:
            obj = self
            for attr in key.split('.')[:-1]:
                if isinstance(obj, AgentList):
                    assert attr.isdigit()
                    obj = obj[int(attr)]
                elif isinstance(obj, AgentDict):
                    obj = obj[attr]
                else:
                    obj = getattr(obj, attr)
            if obj.memory is not None:
                if session_id not in obj.memory.memory_map:
                    obj.memory.create_instance(session_id)
                obj.memory.memory_map[session_id].load(state_dict[key] or [])

    def register_hook(self, hook: Callable):
        handle = RemovableHandle(self._hooks)
        self._hooks[handle.id] = hook
        return handle

    def reset(self, session_id=0, keypath: Optional[str] = None, recursive: bool = False):
        assert not (keypath and recursive), 'keypath and recursive can\'t be used together'
        if keypath:
            keys, agent = keypath.split('.'), self
            for key in keys:
                agents = getattr(agent, '_agents', {})
                if key not in agents:
                    raise KeyError(f'No sub-agent named {key} in {agent}')
                agent = agents[key]
            agent.reset(session_id, recursive=False)
        else:
            if self.memory:
                self.memory.reset(session_id=session_id)
            if recursive:
                for agent in getattr(self, '_agents', {}).values():
                    agent.reset(session_id, recursive=True)

    def __repr__(self):

        def _rcsv_repr(agent, n_indent=1):
            res = agent.__class__.__name__ + (f"(name='{agent.name}')" if agent.name else '')
            modules = [
                f"{n_indent * '  '}({name}): {_rcsv_repr(agent, n_indent + 1)}"
                for name, agent in getattr(agent, '_agents', {}).items()
            ]
            if modules:
                res += '(\n' + '\n'.join(modules) + f'\n{(n_indent - 1) * "  "})'
            elif not res.endswith(')'):
                res += '()'
            return res

        return _rcsv_repr(self)


class AsyncAgentMixin:

    async def __call__(self, *message: AgentMessage, session_id=0, **kwargs) -> AgentMessage:
        message = [AgentMessage(sender='user', content=m) if isinstance(m, str) else copy.deepcopy(m) for m in message]
        for hook in self._hooks.values():
            result = hook.before_agent(self, message, session_id)
            if result:
                message = result
        self.update_memory(message, session_id=session_id)
        response_message = await self.forward(*message, session_id=session_id, **kwargs)
        if not isinstance(response_message, AgentMessage):
            response_message = AgentMessage(sender=self.name, content=response_message)
        self.update_memory(response_message, session_id=session_id)
        response_message = copy.deepcopy(response_message)
        for hook in self._hooks.values():
            result = hook.after_agent(self, response_message, session_id)
            if result:
                response_message = result
        return response_message

    async def forward(self, *message: AgentMessage, session_id=0, **kwargs) -> Union[AgentMessage, str]:
        formatted_messages = self.aggregator.aggregate(
            self.memory.get(session_id), self.name, self.output_format, self.template
        )
        llm_response = await self.llm.chat(formatted_messages, session_id, **kwargs)
        if self.output_format:
            formatted_messages = self.output_format.parse_response(llm_response)
            return AgentMessage(sender=self.name, content=llm_response, formatted=formatted_messages)
        return llm_response


class AsyncAgent(AsyncAgentMixin, Agent):
    """Asynchronous variant of the Agent class"""

    pass


class StreamingAgentMixin:
    """Component that makes agent calling output a streaming response."""

    def __call__(self, *message: AgentMessage, session_id=0, **kwargs) -> Generator[AgentMessage, None, None]:
        message = [AgentMessage(sender='user', content=m) if isinstance(m, str) else copy.deepcopy(m) for m in message]
        for hook in self._hooks.values():
            result = hook.before_agent(self, message, session_id)
            if result:
                message = result
        self.update_memory(message, session_id=session_id)
        response_message = AgentMessage(sender=self.name, content="")
        for response_message in self.forward(*message, session_id=session_id, **kwargs):
            if not isinstance(response_message, AgentMessage):
                model_state, response = response_message
                response_message = AgentMessage(sender=self.name, content=response, stream_state=model_state)
            yield response_message.model_copy()
        self.update_memory(response_message, session_id=session_id)
        response_message = copy.deepcopy(response_message)
        for hook in self._hooks.values():
            result = hook.after_agent(self, response_message, session_id)
            if result:
                response_message = result
        yield response_message

    def forward(
        self, *message: AgentMessage, session_id=0, **kwargs
    ) -> Generator[Union[AgentMessage, Tuple[ModelStatusCode, str]], None, None]:
        formatted_messages = self.aggregator.aggregate(
            self.memory.get(session_id), self.name, self.output_format, self.template
        )
        for model_state, response, *_ in self.llm.stream_chat(formatted_messages, session_id=session_id, **kwargs):
            yield (
                AgentMessage(
                    sender=self.name,
                    content=response,
                    formatted=self.output_format.parse_response(response),
                    stream_state=model_state,
                )
                if self.output_format
                else (model_state, response)
            )


class AsyncStreamingAgentMixin:
    """Component that makes asynchronous agent calling output a streaming response."""

    async def __call__(self, *message: AgentMessage, session_id=0, **kwargs) -> AsyncGenerator[AgentMessage, None]:
        message = [AgentMessage(sender='user', content=m) if isinstance(m, str) else copy.deepcopy(m) for m in message]
        for hook in self._hooks.values():
            result = hook.before_agent(self, message, session_id)
            if result:
                message = result
        self.update_memory(message, session_id=session_id)
        response_message = AgentMessage(sender=self.name, content="")
        async for response_message in self.forward(*message, session_id=session_id, **kwargs):
            if not isinstance(response_message, AgentMessage):
                model_state, response = response_message
                response_message = AgentMessage(sender=self.name, content=response, stream_state=model_state)
            yield response_message.model_copy()
        self.update_memory(response_message, session_id=session_id)
        response_message = copy.deepcopy(response_message)
        for hook in self._hooks.values():
            result = hook.after_agent(self, response_message, session_id)
            if result:
                response_message = result
        yield response_message

    async def forward(
        self, *message: AgentMessage, session_id=0, **kwargs
    ) -> AsyncGenerator[Union[AgentMessage, Tuple[ModelStatusCode, str]], None]:
        formatted_messages = self.aggregator.aggregate(
            self.memory.get(session_id), self.name, self.output_format, self.template
        )
        async for model_state, response, *_ in self.llm.stream_chat(
            formatted_messages, session_id=session_id, **kwargs
        ):
            yield (
                AgentMessage(
                    sender=self.name,
                    content=response,
                    formatted=self.output_format.parse_response(response),
                    stream_state=model_state,
                )
                if self.output_format
                else (model_state, response)
            )


class StreamingAgent(StreamingAgentMixin, Agent):
    """Streaming variant of the Agent class"""

    pass


class AsyncStreamingAgent(AsyncStreamingAgentMixin, Agent):
    """Streaming variant of the AsyncAgent class"""

    pass


class Sequential(Agent):
    """Sequential is an agent container that forwards messages to each agent
    in the order they are added."""

    def __init__(self, *agents: Union[Agent, Iterable], **kwargs):
        super().__init__(**kwargs)
        self._agents = OrderedDict()
        if not agents:
            raise ValueError('At least one agent should be provided')
        if isinstance(agents[0], Iterable) and not isinstance(agents[0], Agent):
            if not agents[0]:
                raise ValueError('At least one agent should be provided')
            agents = agents[0]
        for key, agent in enumerate(agents):
            if isinstance(agents, Mapping):
                key, agent = agent, agents[agent]
            elif isinstance(agent, tuple):
                key, agent = agent
            self.add_agent(key, agent)

    def add_agent(self, name: str, agent: Agent):
        assert isinstance(agent, Agent), f'{type(agent)} is not an Agent subclass'
        self._agents[str(name)] = agent

    def forward(self, *message: AgentMessage, session_id=0, exit_at: Optional[int] = None, **kwargs) -> AgentMessage:
        assert exit_at is None or exit_at >= 0, 'exit_at should be greater than or equal to 0'
        if exit_at is None:
            exit_at = len(self) - 1
        iterator = chain.from_iterable(repeat(self._agents.values()))
        for _ in range(exit_at + 1):
            agent = next(iterator)
            if isinstance(message, AgentMessage):
                message = (message,)
            message = agent(*message, session_id=session_id, **kwargs)
        return message

    def __getitem__(self, key):
        if isinstance(key, int) and key < 0:
            assert key >= -len(self), 'index out of range'
            key = len(self) + key
        return self._agents[str(key)]

    def __len__(self):
        return len(self._agents)


class AsyncSequential(AsyncAgentMixin, Sequential):

    async def forward(
        self, *message: AgentMessage, session_id=0, exit_at: Optional[int] = None, **kwargs
    ) -> AgentMessage:
        assert exit_at is None or exit_at >= 0, 'exit_at should be greater than or equal to 0'
        if exit_at is None:
            exit_at = len(self) - 1
        iterator = chain.from_iterable(repeat(self._agents.values()))
        for _ in range(exit_at + 1):
            agent = next(iterator)
            if isinstance(message, AgentMessage):
                message = (message,)
            message = await agent(*message, session_id=session_id, **kwargs)
        return message


class StreamingSequential(StreamingAgentMixin, Sequential):
    """Streaming variant of the Sequential class"""

    def forward(self, *message: AgentMessage, session_id=0, exit_at: Optional[int] = None, **kwargs):
        assert exit_at is None or exit_at >= 0, 'exit_at should be greater than or equal to 0'
        if exit_at is None:
            exit_at = len(self) - 1
        iterator = chain.from_iterable(repeat(self._agents.values()))
        for _ in range(exit_at + 1):
            agent = next(iterator)
            if isinstance(message, AgentMessage):
                message = (message,)
            for message in agent(*message, session_id=session_id, **kwargs):
                yield message


class AsyncStreamingSequential(AsyncStreamingAgentMixin, Sequential):
    """Streaming variant of the AsyncSequential class"""

    async def forward(self, *message: AgentMessage, session_id=0, exit_at: Optional[int] = None, **kwargs):
        assert exit_at is None or exit_at >= 0, 'exit_at should be greater than or equal to 0'
        if exit_at is None:
            exit_at = len(self) - 1
        iterator = chain.from_iterable(repeat(self._agents.values()))
        for _ in range(exit_at + 1):
            agent = next(iterator)
            if isinstance(message, AgentMessage):
                message = (message,)
            async for message in agent(*message, session_id=session_id, **kwargs):
                yield message


class AgentContainerMixin:

    def __init_subclass__(cls):
        super().__init_subclass__()

        def wrap_api(func):

            @wraps(func)
            def wrapped_func(self, *args, **kwargs):
                data = self.data.copy() if hasattr(self, 'data') else None

                def _backup(d):
                    if d is None:
                        self.data.clear()
                    else:
                        self.data = d

                ret = func(self, *args, **kwargs)
                agents = OrderedDict()
                for k, item in self.data.items() if isinstance(self.data, abc.Mapping) else enumerate(self.data):
                    if isinstance(self.data, abc.Mapping) and not isinstance(k, str):
                        _backup(data)
                        raise KeyError(f'agent name should be a string, got {type(k)}')
                    if isinstance(k, str) and '.' in k:
                        _backup(data)
                        raise KeyError(f'agent name can\'t contain ".", got {k}')
                    if not isinstance(item, Agent):
                        _backup(data)
                        raise TypeError(f'{type(item)} is not an Agent subclass')
                    agents[str(k)] = item
                self._agents = agents
                return ret

            return wrapped_func

        # fmt: off
        for method in [
            'append', 'sort', 'reverse', 'pop', 'clear', 'update',
            'insert', 'extend', 'remove', '__init__', '__setitem__',
            '__delitem__', '__add__', '__iadd__', '__radd__', '__mul__',
            '__imul__', '__rmul__'
        ]:
            if hasattr(cls, method):
                setattr(cls, method, wrap_api(getattr(cls, method)))


class AgentList(Agent, UserList, AgentContainerMixin):

    def __init__(self, agents: Optional[Iterable[Agent]] = None):
        Agent.__init__(self, memory=None)
        UserList.__init__(self, agents)
        self.name = None


class AgentDict(Agent, UserDict, AgentContainerMixin):

    def __init__(self, agents: Optional[Mapping[str, Agent]] = None):
        Agent.__init__(self, memory=None)
        UserDict.__init__(self, agents)
        self.name = None