"""The function tools tht are actually implemented"""

import json
import subprocess

from langchain.agents.load_tools import load_tools
from langchain.tools import BaseTool
from langchain.utilities.bash import BashProcess
#from langchain_experimental.llm_bash.bash import BashProcess
from safeagi.tools.tool_interface import (
    ArgException,
    ArgParameter,
    ArgReturn,
    FunctionTool,
    FunctionToolkit,
)
from safeagi.utils.my_typing import *

from .register import register_toolkit

__ALL__ = ["RealTerminal", "RealPythonInterpreter", "RealWikipedia", "RealHuman"]


class MyBashProcess(BashProcess):
    def _run(self, command: str) -> Tuple[str, int]:
        """
        Runs a command in a subprocess and returns
        the output.

        Args:
            command: The command to run
        """  # noqa: E501
        try:
            output = (
                subprocess.run(
                    command,
                    shell=True,
                    check=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                )
                .stdout.decode()
                .strip()
            )
        except subprocess.CalledProcessError as error:
            if self.return_err_output:
                return error.stdout.decode().strip(), error.returncode

            return str(error).strip(), error.returncode
        if self.strip_newlines:
            output = output.strip()

        return output, 0


#################### Terminal Interpreter ####################
class RealTerminalExecute(FunctionTool):
    name = "TerminalExecute"
    summary = "Execute a terminal command and return the output. This command should follow proper syntax and be supported by the terminal environment."
    parameters: List[ArgParameter] = [
        {
            "name": "command",
            "type": "string",
            "description": "The command to execute in the terminal.",
            "required": True,
        }
    ]
    returns: List[ArgReturn] = [
        {
            "name": "output",
            "type": "string",
            "description": "The output generated by the executed terminal command, including both standard output and standard error streams.",
        },
        {
            "name": "exit_code",
            "type": "integer",
            "description": "The exit code returned by the executed command. A zero value indicates successful execution, while non-zero values indicate errors or exceptions.",
        },
    ]

    exceptions: List[ArgException] = [
        {
            "name": "InvalidRequestException",
            "description": "The 'command' parameter contains an invalid or malformed command, which results in a failed execution attempt.",
        }
    ]

    _tool: BaseTool = MyBashProcess(return_err_output=True)

    def parse_return(self, tool_output: Dict[str, Any]) -> str:
        return json.dumps({"output": tool_output[0], "exit_code": tool_output[1]})

    def _runtool(self, tool_input: Dict[str, Any]) -> Dict[str, Any]:
        return self._tool._run(tool_input["command"])

    def _aruntool(self, tool_input: Dict[str, Any]) -> Dict[str, Any]:
        return self._tool._arun(tool_input["command"])


@register_toolkit()
class RealTerminal(FunctionToolkit):
    name_for_human = "Terminal command executor"
    description_for_human = "Executes commands in a terminal."
    name_for_model = "Terminal"
    description_for_model = "Executes commands in a terminal on the user's local system. Use it to run valid terminal commands for tasks such as file management, system control, and more"
    tool_classes = [RealTerminalExecute]


#################### Python Interpreter ####################


class RealPythonInterpreterExecute(FunctionTool):
    name = "PythonInterpreterExecute"
    summary = "Execute a Python script."
    parameters: List[ArgParameter] = [
        {
            "name": "script",
            "type": "string",
            "description": "The python script to execute.",
            "required": True,
        }
    ]
    returns: List[ArgReturn] = [
        {
            "name": "result",
            "type": "string",
            "description": "The printed output of the script.",
        }
    ]
    exceptions: List[ArgException] = []
    _tool: BaseTool = load_tools(["python_repl"])[0]

    def parse_return(self, tool_output: str) -> str:
        return json.dumps({"result": tool_output})

    def _runtool(self, tool_input: Dict[str, Any]) -> Dict[str, Any]:
        return self._tool._run(tool_input["script"])

    def _aruntool(self, tool_input: Dict[str, Any]) -> Dict[str, Any]:
        return self._tool._arun(tool_input["script"])


@register_toolkit()
class RealPythonInterpreter(FunctionToolkit):
    name_for_human = "Python interpreter"
    description_for_human = "A Python shell."
    name_for_model = "PythonInterpreter"
    description_for_model = "A Python shell. Use it to execute python scripts.  If you want to see the output of a value, you should print it out with `print(...)`."

    tool_classes = [RealPythonInterpreterExecute]


#################### Wikipedia ####################
class RealWikipediaSearch(FunctionTool):
    name = "WikipediaSearch"
    summary = "Query the Wikipedia tool for a given query."
    parameters: List[ArgParameter] = [
        {
            "name": "query",
            "type": "string",
            "description": "The query to search for.",
            "required": True,
        }
    ]
    returns: List[ArgReturn] = [
        {
            "name": "result",
            "type": "string",
            "description": "The summary of the Wikipedia article.",
        }
    ]
    exceptions: List[ArgException] = []
    _tool: BaseTool = load_tools(["wikipedia"])[0]

    def parse_return(self, tool_output: str) -> str:
        return json.dumps({"result": tool_output})

    def _runtool(self, tool_input: Dict[str, Any]) -> Dict[str, Any]:
        return self._tool._run(tool_input["query"])

    def _aruntool(self, tool_input: Dict[str, Any]) -> Dict[str, Any]:
        return self._tool._arun(tool_input["query"])


@register_toolkit()
class RealWikipedia(FunctionToolkit):
    name_for_human = "Wikipedia search tool"
    description_for_human = "Tool for searching through Wikipedia."
    name_for_model = "Wikipedia"
    description_for_model = "Tool for searching through Wikipedia. Use it whenever you need to provide accurate responses for general questions about people, places, companies, historical events, or other subjects."

    tool_classes = [RealWikipediaSearch]


#################### Human ####################


class RealHumanAssistanceQuery(FunctionTool):
    name = "HumanAssistanceQuery"
    summary = "Ask the human a specific question"
    parameters: List[ArgParameter] = [
        {
            "name": "question",
            "type": "string",
            "description": "The question to ask.",
            "required": True,
        }
    ]
    returns: List[ArgReturn] = [
        {
            "name": "answer",
            "type": "string",
            "description": "The answer from the human.",
        }
    ]
    exceptions: List[ArgException] = []

    def parse_return(self, tool_output: str) -> str:
        return json.dumps({"answer": tool_output})

    def _runtool(self, tool_input: Dict[str, Any]) -> Dict[str, Any]:
        print("\n" + tool_input["question"] + "\n")
        return input(tool_input["question"])

    def _aruntool(self, tool_input: Dict[str, Any]) -> Dict[str, Any]:
        return NotImplementedError("Human tool does not support async")


@register_toolkit()
class RealHuman(FunctionToolkit):
    name_for_human = "Human assistance"
    description_for_human = "Seek human assistance or guidance."
    name_for_model = "HumanAssistance"
    description_for_model = "Seek human assistance or guidance. Use it when expert human or user input is necessary, e.g., when you need some human knowledge, user permission, user-specific information."
    tool_classes = [RealHumanAssistanceQuery]
