from typing import Any

from pydantic_ai import RunContext  # type: ignore[import-not-found]
from pydantic_ai.tools import Tool as PydanticTool  # type: ignore[import-not-found]
from pydantic_ai.tools import ToolDefinition  # type: ignore[import-not-found]

from blaxel.core.tools import Tool
from blaxel.core.tools import bl_tools as bl_tools_core


def _clean_schema_for_openai(schema: dict[str, Any]) -> dict[str, Any]:
    """Clean JSON schema to be compatible with OpenAI function calling.

    OpenAI requires object schemas to have a 'properties' field, even if empty.
    """
    if not isinstance(schema, dict):
        return schema

    cleaned = schema.copy()

    if cleaned.get("type") == "object":
        if "properties" not in cleaned:
            cleaned["properties"] = {}
        if "required" not in cleaned:
            cleaned["required"] = []

    if "additionalProperties" in cleaned:
        del cleaned["additionalProperties"]
    if "$schema" in cleaned:
        del cleaned["$schema"]

    if "properties" in cleaned:
        cleaned["properties"] = {
            k: _clean_schema_for_openai(v) for k, v in cleaned["properties"].items()
        }
    if "items" in cleaned and isinstance(cleaned["items"], dict):
        cleaned["items"] = _clean_schema_for_openai(cleaned["items"])

    return cleaned


def get_pydantic_tool(tool: Tool) -> PydanticTool:
    """
    Converts a custom Tool object into a Pydantic AI Tool object.

    Uses the attributes defined in the custom Tool (name, description, input_schema)
    by defining a dynamic 'prepare' function that modifies the ToolDefinition
    generated by Pydantic AI.
    """
    # Select the function to use, prioritizing the async version
    func = tool.coroutine if tool.coroutine else tool.sync_coroutine
    if not func:
        raise ValueError(
            f"Tool '{tool.name}' must have either a coroutine or sync_coroutine defined."
        )

    # Define the prepare function dynamically to capture the 'tool' object
    # Assuming RunContext and ToolDefinition types based on the pydantic_ai example
    async def prepare_tool(ctx: RunContext, tool_def: ToolDefinition) -> ToolDefinition | None:
        """Dynamically prepares the ToolDefinition using the custom Tool's attributes."""
        tool_def.name = tool.name  # Override inferred name
        tool_def.description = tool.description  # Override inferred description
        tool_def.parameters_json_schema = _clean_schema_for_openai(tool.input_schema)
        return tool_def

    async def pydantic_function(**kwargs):
        return await func(**kwargs)

    # Create the Pydantic AI Tool, passing the function and the prepare hook
    return PydanticTool(
        pydantic_function,
        name=tool.name,
        description=tool.description,
        prepare=prepare_tool,
        takes_ctx=False,
    )


async def bl_tools(tools_names: list[str], **kwargs) -> list[PydanticTool]:
    tools = bl_tools_core(tools_names, **kwargs)
    await tools.initialize()
    return [get_pydantic_tool(tool) for tool in tools.get_tools()]
