# type: ignore
from typing import Any, Dict, Type
from pydantic import BaseModel, create_model, Field
from mcp import ClientSession
from .base import BaseTool


class MCPTool(BaseTool):
    """
    A proxy tool that forwards calls to an MCP Client Session.
    """

    def __init__(
        self,
        name: str,
        description: str,
        input_schema: Dict[str, Any],
        session: ClientSession,
    ):
        self.name = name
        self.description = description
        self._session = session
        self._input_schema = input_schema
        self.args_schema = self._create_args_schema(input_schema)

    def _create_args_schema(self, schema: Dict[str, Any]) -> Type[BaseModel]:
        """
        Dynamically create a Pydantic model from JSON Schema.
        """
        properties = schema.get("properties", {})
        required = schema.get("required", [])

        fields = {}
        for prop_name, prop_def in properties.items():
            prop_type = str  # Default to string

            t = prop_def.get("type")
            if t == "integer":
                prop_type = int
            elif t == "number":
                prop_type = float
            elif t == "boolean":
                prop_type = bool
            elif t == "array":
                prop_type = list
            elif t == "object":
                prop_type = dict

            description = prop_def.get("description", "")

            if prop_name in required:
                fields[prop_name] = (prop_type, Field(..., description=description))
            else:
                from typing import Optional as Opt

                fields[prop_name] = (
                    Opt[prop_type],
                    Field(None, description=description),
                )

        return create_model(f"{self.name}Args", **fields)

    async def run(self, **kwargs) -> Any:
        """
        Execute the tool via MCP protocol.
        """
        # Call the tool on remote server
        result = await self._session.call_tool(self.name, arguments=kwargs)

        # Format result
        output = []
        if hasattr(result, "content"):
            for item in result.content:
                if hasattr(item, "text"):
                    output.append(item.text)
                else:
                    output.append(str(item))
        else:
            output.append(str(result))

        return "\n".join(output)
