import json
import re
from json import JSONDecodeError
from typing import List, Tuple, Any, Union
from dataclasses import dataclass

from langchain.callbacks.manager import Callbacks
from langchain.agents import (
    OpenAIFunctionsAgent,
)
from langchain.schema import (
    AgentAction,
    AgentFinish,
    BaseMessage,
    OutputParserException
)
from langchain.schema.messages import (
    AIMessage,
    FunctionMessage
)
from langchain.tools import BaseTool, StructuredTool
# from langchain.tools.convert_to_openai import FunctionDescription
from langchain.utils.openai_functions import FunctionDescription


@dataclass
class _FunctionsAgentAction(AgentAction):
    """Add message_log to AgentAction class for the _FunctionAgentAction
    """
    message_log: List[BaseMessage]

    def __init__(
            self, tool: str, tool_input: Union[str, dict], log: str, **kwargs: Any
    ):
        """Override init to support instantiation by position for backward compat."""
        super().__init__(tool=tool, tool_input=tool_input, log=log, **kwargs)


def _convert_agent_action_to_messages(
        agent_action: AgentAction, observation: str
) -> List[BaseMessage]:
    """Convert an agents action to a message.

    This code is used to reconstruct the original AI message from the agents action.

    Args:
        agent_action: Agent action to convert.

    Returns:
        AIMessage that corresponds to the original tools invocation.
    """
    if isinstance(agent_action, _FunctionsAgentAction):
        return agent_action.message_log + [
            _create_function_message(agent_action, observation)
        ]
    else:
        return [AIMessage(content=agent_action.log)]


def _create_function_message(
        agent_action: AgentAction, observation: str
) -> FunctionMessage:
    """Convert agents action and observation into a function message.
    Args:
        agent_action: the tools invocation request from the agents
        observation: the result of the tools invocation
    Returns:
        FunctionMessage that corresponds to the original tools invocation
    """
    if not isinstance(observation, str):
        try:
            content = json.dumps(observation, ensure_ascii=False)
        except Exception:
            content = str(observation)
    else:
        content = observation
    return FunctionMessage(
        name=agent_action.tool,
        content=content,
    )


def _format_intermediate_steps(
        intermediate_steps: List[Tuple[AgentAction, str]],
) -> List[BaseMessage]:
    """Format intermediate steps.
    Args:
        intermediate_steps: Steps the LLM has taken to date, along with observations
    Returns:
        list of messages to send to the LLM for the next prediction
    """
    messages = []

    for intermediate_step in intermediate_steps:
        agent_action, observation = intermediate_step
        messages.extend(_convert_agent_action_to_messages(agent_action, observation))

    return messages


def _format_tool_to_openai_function(tool: BaseTool) -> FunctionDescription:
    """Format tools into the OpenAI function API."""
    parameters = {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": tool.param_description if hasattr(tool, 'param_description') else "",
            },
        },
        "required": ["query"],
    }

    return {
        "name": tool.name,
        "description": tool.description,
        "parameters": parameters,
    }


def _parse_ai_message(message: BaseMessage) -> Union[AgentAction, AgentFinish]:
    if not isinstance(message, AIMessage):
        raise TypeError(f"Expected an AI message but got {type(message)}")

    function_call = message.additional_kwargs.get("function_call", {})

    if function_call:
        function_call = message.additional_kwargs["function_call"]
        function_name = function_call["name"]
        try:
            _tool_input = json.loads(function_call["arguments"])
        except JSONDecodeError:
            raise OutputParserException(
                f"Could not parse tools input: {function_call} because"
                f"the `arguments` is not valid JSON."
            )

        if "query" in _tool_input:
            tool_input = _tool_input["query"]
        else:
            tool_input = _tool_input

        return _FunctionsAgentAction(
            tool=function_name,
            tool_input=tool_input,
            log=f"\nInvoking: `{function_name}` with `{tool_input}`\n",
            message_log=[message]
        )

    # pattern = r'\((.*?)\)'
    # matches = re.findall(pattern, message.content)
    # result = []
    #
    # for match in matches:
    #     result.append(match)
    #
    # if result:
    #     output = result
    # else:
    #     output = message.content

    return AgentFinish(return_values={"output": message.content}, log=message.content)


class ConversationalFunctionsAgent(OpenAIFunctionsAgent):
    @property
    def functions(self) -> List[dict]:
        return [dict(_format_tool_to_openai_function(t)) for t in self.tools]

    def plan(self,
             intermediate_steps: List[Tuple[AgentAction, str]],
             callbacks: Callbacks = None,
             **kwargs: Any
             ) -> Union[AgentAction, AgentFinish]:
        """Decide how agents should move after receiving an input. The difference between
        OpenAIFunctionsAgent lies in the '_parse_ai_message' function. We add an OutputParser
        into it.

        Args:
            intermediate_steps: Steps the LLM has taken to date, along with observations
            **kwargs: User inputs.
            **kwargs: Including user's input string

        Returns:
            Action specifying what tools to use.
        """
        agent_scratchpad: List[BaseMessage] = _format_intermediate_steps(intermediate_steps)
        selected_inputs = {
            k: kwargs[k] for k in self.prompt.input_variables if k != "agent_scratchpad"
        }
        full_inputs = dict(**selected_inputs, agent_scratchpad=agent_scratchpad)
        prompt = self.prompt.format_prompt(**full_inputs)
        messages: List[BaseMessage] = prompt.to_messages()
        predicted_message = self.llm.predict_messages(
            messages, functions=self.functions, callbacks=callbacks
        )
        agent_decision = _parse_ai_message(predicted_message)
        return agent_decision
