import asyncio
import inspect
import textwrap
from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, Union

from pydantic import BaseModel, Field, JsonValue, ValidationError
from agentia.models.base import File
from agentia.tools.mcp import MCP
import agentia.models as models
from agentia.utils.decorators import ToolFuncParam
from inspect import Parameter
from pydantic.type_adapter import TypeAdapter

NAME_TAG = "agentia_tool_name"
DISPLAY_NAME_TAG = "agentia_tool_display_name"
IS_TOOL_TAG = "agentia_tool_is_tool"
DESCRIPTION_TAG = "agentia_tool_description"
METADATA_TAG = "agentia_tool_metadata"

if TYPE_CHECKING:
    from agentia.plugins import Plugin
    from agentia.agent import Agent


class _BaseTool:
    def __init__(
        self,
        name: str,
        schema: Any,
        display_name: str | None = None,
        description: str | None = None,
        metadata: Any | None = None,
    ):
        self.name = name
        self.display_name = display_name
        self.description = description
        self.schema = schema
        self.metadata = metadata


class _PythonFunctionTool(_BaseTool):
    def __init__(self, f: Callable[..., Any], plugin: Optional["Plugin"] = None):
        if plugin:
            name = getattr(f, NAME_TAG, None) or f"{plugin.name()}__{f.__name__}"
            display_name: str = (
                getattr(f, DISPLAY_NAME_TAG, None) or f"{plugin.name()}@{f.__name__}"
            )
        else:
            name = getattr(f, NAME_TAG, f.__name__)
            display_name = getattr(f, DISPLAY_NAME_TAG, name)
        description: str = textwrap.dedent(getattr(f, DESCRIPTION_TAG, f.__doc__) or "")
        metadata = getattr(f, METADATA_TAG, None)
        params = [
            ToolFuncParam(p, f.__name__)
            for p in inspect.signature(f).parameters.values()
            if p.name != "self"
        ]
        schema = _PythonFunctionTool.get_json_schema(
            name=name, description=description, fparams=params
        )
        super().__init__(
            name=name,
            schema=schema,
            display_name=display_name,
            description=description,
            metadata=metadata,
        )
        self.func = f
        self.plugin = plugin
        self.params = params

    @staticmethod
    def get_json_schema(
        name: str, description: str | None, fparams: list[ToolFuncParam]
    ) -> Any:
        params: Any = {"type": "object", "properties": {}, "required": []}
        for p in fparams:
            if p.required:
                params["required"].append(p.name)
            params["properties"][p.name] = p.schema
        return {
            "name": name,
            "description": description or "",
            "parameters": params,
        }

    def process_json_args(self, agent: "Agent", args: Any):
        from agentia.agent import Agent

        p_args: list[Any] = []
        kw_args: dict[str, Any] = {}

        def get_value(p: ToolFuncParam, name: str, default: Any, is_model: bool):
            if name not in args:
                return default
            if is_model:
                return p.type(**args[name])
            else:
                v = args[name]
                try:
                    return TypeAdapter(p.type).validate_python(v)
                except ValidationError as e:
                    raise ValueError(
                        f"Error parsing argument '{name}': {e}. Please retry with a valid value."
                    ) from e

        assert self.params is not None
        for p in self.params:
            is_model = inspect.isclass(p.type) and issubclass(p.type, BaseModel)
            match p.param.kind:
                case Parameter.POSITIONAL_ONLY if p.param.annotation == Agent:
                    if agent is not None:
                        p_args.append(agent)
                case Parameter.POSITIONAL_ONLY:
                    default = p.default if p.default != Parameter.empty else None
                    p_args.append(get_value(p, p.name, default, is_model))
                case Parameter.POSITIONAL_OR_KEYWORD | Parameter.KEYWORD_ONLY if (
                    p.param.annotation == Agent
                ):
                    kw_args[p.name] = agent
                case Parameter.POSITIONAL_OR_KEYWORD | Parameter.KEYWORD_ONLY:
                    default = p.default if p.default != Parameter.empty else None
                    kw_args[p.name] = get_value(p, p.name, default, is_model)
                case other:
                    raise ValueError(f"{other} is not supported")
        return p_args, kw_args


class _MCPTool(_BaseTool):
    def __init__(self, name: str, schema: Any, server: MCP):
        super().__init__(name=name, schema=schema)
        self.server = server


class ProviderTool(BaseModel):
    name: str
    args: dict[str, JsonValue] | None = None


type Tool = Plugin | Callable[..., Any] | ProviderTool | MCP
type Tools = Sequence[Tool]


class ToolResult(BaseModel):
    output: JsonValue
    files: list[File] = Field(default_factory=list)


class ToolSet:
    def __init__(self, tools: Tools, agent: "Agent"):
        from agentia.plugins import Plugin

        self.all_tools = tools
        self.plugins: dict[str, Plugin] = {}
        self.mcp_servers: dict[str, MCP] = {}
        self.agent = agent
        self.provider_tools: dict[str, ProviderTool] = {}
        self.__tools: dict[str, _BaseTool] = {}
        for t in tools:
            if isinstance(t, Plugin):
                self.__add_plugin(t)
            elif isinstance(t, MCP):
                self.__add_mcp_server(t)
            elif isinstance(t, ProviderTool):
                self.provider_tools[t.name] = t
            else:
                assert inspect.isfunction(t), (
                    "Expected a function, MCP server object, Plugin, or ProviderTool"
                )
                self.__add_function(t)
        self.__initialized = False

    async def init(self):
        if self.__initialized:
            return
        self.__initialized = True
        for plugin in self.plugins.values():
            try:
                plugin._agent = self.agent
                await plugin.init()
            except Exception as e:
                from . import PluginInitError

                raise PluginInitError(plugin.id(), e) from e
        for server in self.mcp_servers.values():
            await server.init(agent=self.agent)
            tools = server.get_tools()
            for tool in tools:
                self.__tools[tool.name] = tool

    def add(self, t: Union[Callable[..., Any], MCP, "Plugin"]) -> None:
        from agentia.plugins import Plugin

        if isinstance(t, Plugin):
            self.__add_plugin(t)
        elif isinstance(t, MCP):
            self.__add_mcp_server(t)
        else:
            assert inspect.isfunction(t), (
                "Expected a function, MCP server object, or Plugin"
            )
            self.__add_function(t)

    def __add_function(self, f: Callable[..., Any], plugin: Optional["Plugin"] = None):
        t = _PythonFunctionTool(f=f, plugin=plugin)
        self.__tools[t.name] = t

    def __add_plugin(self, p: "Plugin"):
        # Add all functions from the plugin
        for _n, method in inspect.getmembers(p, predicate=inspect.ismethod):
            if not getattr(method, IS_TOOL_TAG, False):
                continue
            self.__add_function(method, plugin=p)
        # Add the plugin to the list of plugins
        self.plugins[p.id()] = p

    def __add_mcp_server(self, server: MCP):
        self.mcp_servers[server.name] = server

    def get_plugin[T: "Plugin"](self, type: type[T]) -> Optional[T]:
        for p in self.plugins.values():
            if isinstance(p, type):
                return p
        return None

    def get_instructions(self) -> str | None:
        instructions = []
        for plugin in self.plugins.values():
            ins = plugin.get_instructions()
            if ins:
                if s := ins.strip():
                    instructions.append(s)
        if not instructions:
            return None
        return "\n\n".join(instructions)

    def is_empty(self) -> bool:
        return len(self.__tools) == 0

    def get_schema(self) -> Sequence[models.ToolSchema]:
        functions = [
            models.FunctionTool(
                name=v.name,
                description=v.description or "",
                input_schema=v.schema["parameters"],
            )
            for (k, v) in self.__tools.items()
        ]
        provider_tools = [
            models.ProviderDefinedTool(id=t.name, name=t.name, args=t.args or {})
            for t in self.provider_tools.values()
        ]
        return functions + provider_tools

    async def __run_python_tool(
        self,
        id: str,
        tool: _PythonFunctionTool,
        agent: "Agent",
        args: dict[str, JsonValue],
    ) -> models.ToolCallResponse:
        p_args, kw_args = tool.process_json_args(agent, args)
        output = tool.func(*p_args, **kw_args)  # type: ignore
        if inspect.isawaitable(output):
            output = await output
        files: list[File] = []
        if isinstance(output, ToolResult):
            files = output.files
            output = output.output
        return models.ToolCallResponse(
            tool_call_id=id,
            tool_name=tool.name,
            input=args,
            output=output,
            output_files=files,
        )

    async def __run_mcp_tool(
        self, id: str, tool: _MCPTool, args: dict[str, JsonValue]
    ) -> models.ToolCallResponse:
        result = await tool.server.run(tool.name, args)
        return models.ToolCallResponse(
            tool_call_id=id, tool_name=tool.name, input=args, output=result
        )

    async def __run_one(
        self, agent: "Agent", c: models.ToolCall
    ) -> models.ToolCallResponse:
        tool = self.__tools.get(c.tool_name, None)
        if not tool:
            raise ValueError(f"Tool {c.tool_name} not found")
        try:
            if isinstance(tool, _PythonFunctionTool):
                return await self.__run_python_tool(
                    id=c.tool_call_id,
                    tool=tool,
                    agent=agent,
                    args=c.input,
                )
            elif isinstance(tool, _MCPTool):
                return await self.__run_mcp_tool(
                    id=c.tool_call_id,
                    tool=tool,
                    args=c.input,
                )
            else:
                raise ValueError(f"Tool {c.tool_name} is of unknown type")
        except Exception as e:
            output: JsonValue = {"error": str(e)}
            r = models.ToolCallResponse(
                tool_call_id=c.tool_call_id,
                tool_name=c.tool_name,
                input=c.input,
                output=output,
            )
            return r

    async def run(
        self, agent: "Agent", tool_calls: list[models.ToolCall], parallel: bool
    ) -> list[models.ToolCallResponse]:
        if parallel:
            results = await asyncio.gather(
                *[self.__run_one(agent, c) for c in tool_calls]
            )
        else:
            results: list[models.ToolCallResponse] = []
            for c in tool_calls:
                results.append(await self.__run_one(agent, c))
        return results
