from typing import Dict, List, Optional, Union

from lagent.agents.aggregator.default_aggregator import DefaultAggregator
from lagent.memory.base_memory import Memory
from lagent.prompts.parsers.tool_parser import MixedToolParser, ToolParser, ToolStatusCode


class InternLMToolAggregator(DefaultAggregator):

    def __init__(self,
                 environment_role='environment',
                 environment_begin='',
                 environment_end='',
                 user_names: Optional[List[str]] = None,
                 few_shot: Optional[List[List[dict]]] = None):
        self.environment_role = environment_role
        self.environment_begin = environment_begin
        self.environment_end = environment_end
        self.user_names = user_names or ['user']
        self.few_shot = few_shot or []

    def aggregate(self,
                  messages: Memory,
                  name: str,
                  parser: Union[ToolParser, MixedToolParser],
                  system_instruction: str = None) -> List[Dict[str, str]]:
        _message = []
        messages = messages.get_memory()
        if system_instruction:
            _message.extend(
                self.aggregate_system_intruction(system_instruction))
        tool_instruction = parser.format_instruction()
        if tool_instruction:
            if isinstance(tool_instruction, str):
                tool_instruction = dict(
                    role='system', content=tool_instruction)
                if parser.tool_type:
                    tool_instruction['name'] = parser.tool_type
            if isinstance(tool_instruction, dict):
                tool_instruction = [tool_instruction]
            _message.extend(tool_instruction)

        for shot in self.few_shot:
            i = 0
            while i < len(shot):
                msg = shot[i]
                if msg['role'] in ['assistant', 'user', 'system']:
                    _message.append(msg)
                elif msg['role'] == self.environment_role:
                    if not msg['content'].startswith(self.environment_begin):
                        msg['content'] = self.environment_begin + msg['content']
                    if not msg['content'].endswith(self.environment_end):
                        msg['content'] += self.environment_end
                    _message.append(msg)
                elif msg['role'] in ['thought', 'language']:
                    if i < len(shot) - 1 and shot[i + 1]['role'] == 'tool':
                        _message.append(
                            dict(
                                role='assistant',
                                content=parser.format_response(
                                    dict(
                                        tool_type=shot[i + 1]['name'],
                                        thought=msg['content'],
                                        action=shot[i + 1]['content'],
                                        status=None))))
                        i += 1
                    else:
                        _message.append(
                            dict(
                                role='assistant',
                                content=parser.format_response(
                                    dict(
                                        tool_type=None,
                                        thought=msg['content'],
                                        action=None,
                                        status=None))))
                else:
                    raise KeyError(f'Unkown role: {msg["role"]}')
                i += 1

        tool_type = None
        for message in messages:
            if message.sender == name:
                if isinstance(message.formatted, dict):
                    parsed = message.formatted
                    if parsed['status'] == ToolStatusCode.PARSING_ERROR:
                        continue
                    _message.append(
                        dict(
                            role='assistant',
                            content=parser.format_response(parsed)))
                    tool_type = parsed['tool_type']
                else:
                    _message.append(
                        dict(role='assistant', content=str(message.content)))
            elif message.sender in self.user_names:
                _message.append(dict(role='user', content=message.content))
            else:
                msg = dict(
                    role=self.environment_role,
                    content=self.environment_begin + str(message.content) +
                    self.environment_end)
                if tool_type:
                    msg['name'] = tool_type
                _message.append(msg)
        return _message