from typing import List, Tuple

from mmengine.registry import Registry

REGISTRY = Registry('helper')


class LangchainAgent:
    """Agent wrapper for Langchain.

    https://github.com/langchain-ai/langchain.
    """

    def __init__(self, agent_type, llm, tools) -> None:
        from langchain.agents import initialize_agent, load_tools

        llm = REGISTRY.build(llm)
        tools = load_tools(tools, llm=llm)
        self.agent = initialize_agent(tools,
                                      llm,
                                      agent=agent_type,
                                      return_intermediate_steps=True)

    def chat(self, user_input, ice=None) -> Tuple[str, List[dict]]:
        from langchain.schema import AgentAction
        try:
            generation = self.agent(user_input)
            answer = generation['output']
            steps = []
            for step in generation['intermediate_steps']:
                action: AgentAction = step[0]
                steps.append(
                    dict(
                        type=action.tool,
                        args=action.tool_input,
                        result=step[1],
                        thought=action.log,
                        state=0,
                        errmsg=None,
                    ))
        except Exception as e:
            answer = None
            steps = [
                dict(
                    type='InvalidAction',
                    args={},
                    result=None,
                    thought=None,
                    state=-1002,
                    errmsg=str(e),
                )
            ]
        return answer, steps