import json
import logging
from json import JSONDecodeError
from typing import List, Type, T, Dict, Callable
from pydantic import BaseModel
from enum import Enum
from dataclasses import dataclass

from .exceptions import MaxToolErrorIter, ToolError, UnknownResponseFromLLM
from .messages import OpenAiToolCallingMessage, OpenAIFunctionCallingMessage, OllamaToolCallingMessage
from .tool import Tool
from .base_tool_caller import BaseToolCaller
from .utils import AgentType


class Status(Enum):
    CALL_ERROR = 1
    CUSTOM_ERROR = 2
    OKAY = 3


@dataclass
class ErrorCount:
    call_error: int
    custom_error: int


class ToolErrorsTracking:
    """
        Simple class to track errors in tools during Task execution. There's 2 type of errors that can occur inside a tool.
        When a tool signals an error then we decrement the error counter which when reaching 0 triggers an exception.
    """

    def __init__(self, tools: List[Tool]):
        # Map of tool name to its error count
        self.session_status = {}
        # We initialize the error count for each tool based on the max errors defined in the tool constructor
        for tool in tools:
            self.session_status[tool.tool_name] = ErrorCount(call_error=tool.max_call_error, custom_error=tool.max_custom_error)

    def log_new_error(self, tool: Tool, error_type: Status):
        """Logs a new error for a given tool and error type by decrementing associated counter until it reaches 0."""
        if error_type == Status.CALL_ERROR:
            self.session_status[tool.tool_name].call_error -= 1
        elif error_type == Status.CUSTOM_ERROR:
            self.session_status[tool.tool_name].custom_error -= 1

        if self.session_status[tool.tool_name].call_error < 0:
            raise MaxToolErrorIter(
                f"Too many errors were raise by the tool '{tool.tool_name}'. Stopping after {tool.max_custom_error} errors. You can change the maximum errors a tool can raise in the Tool constructor with @max_custom_error.")
        if self.session_status[tool.tool_name].custom_error < 0:
            raise MaxToolErrorIter(
                f"Too many errors occurred while trying to call the python function by Yacana (tool name: {tool.tool_name}). Stopping after {tool.max_call_error} errors. You can change the maximum call error in the Tool constructor with @max_call_error.")


class OpenAiToolCaller(BaseToolCaller):

    def propose_tool(self, task: str, tools: List[Tool], json_output: bool, structured_output: Type[BaseModel] | None, medias: List[str] | None, streaming_callback: Callable | None = None, task_runtime_config: Dict | None = None, tags: List[str] | None = None):
        self.propose_tools(task, tools, json_output, structured_output, medias, streaming_callback, task_runtime_config, tags)

    def propose_tools(self, task: str, tools: List[Tool], json_output: bool, structured_output: Type[BaseModel] | None, medias: List[str] | None, streaming_callback: Callable | None = None, task_runtime_config: Dict | None = None, tags: List[str] | None = None):
        error_tracking = ToolErrorsTracking(tools)
        retry = False

        while True:
            if not retry:
                self.agent._chat(self.agent.history, task, medias=medias, json_output=json_output, structured_output=structured_output, tools=tools)
            retry = False
            if isinstance(self.agent.history.get_last_message(), OpenAIFunctionCallingMessage):
                for tool_call in self.agent.history.get_last_message().tool_calls:
                    tool = next((tool for tool in tools if tool.tool_name == tool_call.name), None)
                    if tool is None:
                        raise ValueError(f"Tool {tool_call.name} not found in tools list")
                    logging.debug("Found tool: %s", tool.tool_name)
                    func_arguments: dict = {}
                    try:
                        if isinstance(tool_call.arguments, dict):
                            func_arguments = tool_call.arguments
                        else:
                            func_arguments: dict = json.loads(tool_call.arguments)
                            if not isinstance(func_arguments, dict): # In case of multi JSON encoding layers, arguments mays still not be a dict
                                raise ValueError(f"Still not a valid dict after doing json.loads on arguments.")
                    except Exception as e:
                        raise UnknownResponseFromLLM(f"Tool {tool_call.name} arguments are not a valid dict object: `{tool_call.arguments}`. Type: {type(tool_call.arguments)}. Original error: {e}.")
                    tool_output, status = self._call_openai_tool(tool, func_arguments)
                    if self.agent._agent_type == AgentType.OPENAI:
                        self.agent.history.add_message(OpenAiToolCallingMessage(tool_output, tool_call.call_id, tags=self.agent._tags))
                    elif self.agent._agent_type == AgentType.OLLAMA:
                        self.agent.history.add_message(OllamaToolCallingMessage(tool_output, tool_call.name, tags=self.agent._tags))

                    if status in (Status.CUSTOM_ERROR, Status.CALL_ERROR):
                        error_tracking.log_new_error(tool, status)
                        task = "Retry tool calling using newly acquired knowledge. Follow Openai's tool calling JSON."
                        retry = True

                if retry:
                    continue
                logging.info(f"[PROMPT][To: {self.agent.name}]: Retrying with original task and tools answer: '{task}'")
                self.agent._chat(self.agent.history, None, medias=medias, json_output=json_output, structured_output=structured_output, streaming_callback=streaming_callback, tools=tools)
                # If the LLM doesn't like the tool's output it can answer with a new function call. If it does, we use the retry mechanism to restart these steps.
                if isinstance(self.agent.history.get_last_message(), OpenAIFunctionCallingMessage):
                    retry = True
                    continue
                else:
                    break

    def _call_openai_tool(self, tool: Tool, function_args: Dict) -> (str, Status):
        """
        Executes a tool call and handles any errors that occur.

        Parameters
        ----------
        tool : Tool
            The tool to execute.
        function_args : Dict
            The arguments to pass to the tool function.

        Returns
        -------
        str
            The output from the tool execution.
        Status
            The status of the tool execution.
        """
        tool_output: str = ""
        status: Status = Status.OKAY

        try:
            if tool.is_mcp:
                tool_output: str = tool.function_ref(tool_name=tool.tool_name, arguments=function_args)
            else:
                tool_output: str = tool.function_ref(**function_args)
            if tool_output is None:
                tool_output = f"Tool returned 'None' so LLM won't be asked to reflect on the tool result."
            else:
                tool_output = str(tool_output)
            logging.info(f"[TOOL_RESPONSE][{tool.tool_name}]: {tool_output}\n")
        except (ToolError, TypeError, JSONDecodeError) as e:
            if type(e) is ToolError or type(e) is JSONDecodeError:
                logging.warning(f"Tool '{tool.tool_name}' raised an error\n")
                status = Status.CUSTOM_ERROR
                tool_output = e.message
            elif type(e) is TypeError:
                logging.warning(f"Yacana failed to call tool '{tool.tool_name}' correctly based on the LLM output\n")
                tool_output = str(e)
                status = Status.CALL_ERROR
            return tool_output, status
        return tool_output, Status.OKAY
