from pathlib import Path
from typing import Any, Optional, Union
from pydantic import AnyUrl
from fastmcp.client import Client
from fastmcp.client.transports import (
    StreamableHttpTransport,
    PythonStdioTransport,
    ClientTransport,
)
from mcp.types import *
from fastmcp.exceptions import ClientError
import asyncio


class MCPClient:
    def __init__(self, name: str, client: Client):
        self.name = name
        self.client = client

    def __repr__(self) -> str:
        return f"<ClientWrapper(name={self.name})>"

    @classmethod
    def from_http(cls, name: str, url: str | AnyUrl) -> "MCPClient":
        transport = StreamableHttpTransport(url=str(url))
        client = Client(transport=transport)
        return cls(name=name, client=client)

    @classmethod
    def from_stdio(
        cls,
        name: str,
        script_path: Union[str, Path],
        args: Optional[list[str]] = None,
        cwd: Optional[Union[str, Path]] = None,
        env: Optional[dict[str, str]] = None,
    ) -> "MCPClient":
        transport = PythonStdioTransport(
            script_path=script_path, args=args, cwd=cwd, env=env
        )
        client = Client(transport=transport)
        return cls(name=name, client=client)

    def is_connected(self) -> bool:
        return self.client.is_connected()

    async def __aenter__(self):
        await self.client.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.client.__aexit__(exc_type, exc_val, exc_tb)

    # Delegación MCP

    async def call_tool(self, name: str, arguments: dict[str, Any] | None = None) -> list[TextContent | ImageContent | EmbeddedResource]:
        try:
            return await asyncio.wait_for(
                self.client.call_tool(name, arguments), timeout=10
            )
        except asyncio.TimeoutError:
            raise RuntimeError(f"Timeout al ejecutar herramienta '{name}'")

    async def get_prompt(
        self, name: str, arguments: dict[str, str] | None = None
    ) -> GetPromptResult:
        return await self.client.get_prompt(name, arguments)

    async def list_tools(self) -> list[Tool]:
        return await self.client.list_tools()

    async def list_prompts(self) -> list[Prompt]:
        return await self.client.list_prompts()

    async def list_resources(self) -> list[Resource]:
        return await self.client.list_resources()

    async def list_resource_templates(self) -> list[ResourceTemplate]:
        return await self.client.list_resource_templates()

    async def read_resource(
        self, uri: AnyUrl | str
    ) -> list[TextResourceContents | BlobResourceContents]:
        return await self.client.read_resource(uri)

    async def complete(
        self,
        ref: ResourceReference | PromptReference,
        argument: dict[str, str],
    ) -> Completion:
        return await self.client.complete(ref, argument)

    async def ping(self) -> bool:
        return await self.client.ping()

    async def set_logging_level(self, level: LoggingLevel) -> None:
        return await self.client.set_logging_level(level)

    async def send_roots_list_changed(self) -> None:
        return await self.client.send_roots_list_changed()