################################################################################
#  Licensed to the Apache Software Foundation (ASF) under one
#  or more contributor license agreements.  See the NOTICE file
#  distributed with this work for additional information
#  regarding copyright ownership.  The ASF licenses this file
#  to you under the Apache License, Version 2.0 (the
#  "License"); you may not use this file except in compliance
#  with the License.  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
import asyncio
import base64
from abc import ABC
from contextlib import AsyncExitStack, asynccontextmanager
from datetime import timedelta
from typing import Any, AsyncIterator, Dict, List
from urllib.parse import urlparse

import cloudpickle
import httpx
from mcp.client.session import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from mcp.types import PromptArgument, TextContent
from pydantic import (
    ConfigDict,
    Field,
    field_serializer,
    model_validator,
)
from typing_extensions import override

from flink_agents.api.chat_message import ChatMessage, MessageRole
from flink_agents.api.prompts.prompt import Prompt
from flink_agents.api.resource import Resource, ResourceType
from flink_agents.api.tools.tool import Tool, ToolMetadata, ToolType
from flink_agents.integrations.mcp.utils import extract_mcp_content_item


class MCPTool(Tool):
    """MCP tool definition that can be called directly.

    This represents a single tool from an MCP server.
    """

    mcp_server: "MCPServer" = Field(default=None)

    @classmethod
    @override
    def tool_type(cls) -> ToolType:
        return ToolType.MCP

    @override
    def call(self, *args: Any, **kwargs: Any) -> Any:
        """Call the MCP tool with the given arguments."""
        if self.mcp_server is None:
            msg = "MCP tool call requires a reference to the MCP server"
            raise ValueError(msg)

        return asyncio.run(
            self.mcp_server.call_tool_async(self.metadata.name, *args, **kwargs)
        )

    @override
    def close(self) -> None:
        if self.mcp_server is not None:
            try:
                self.mcp_server.close()
            finally:
                self.mcp_server = None


class MCPPrompt(Prompt):
    """MCP prompt definition that extends the base Prompt class.

    This represents a prompt managed by an MCP server.
    """

    name: str
    title: str | None
    description: str | None = None
    prompt_arguments: list[PromptArgument] = Field(default_factory=list)
    mcp_server: "MCPServer" = Field(default=None)

    def _check_arguments(self, **kwargs: str) -> Dict[str, str]:
        if self.mcp_server is None:
            msg = "MCP prompt requires a reference to the MCP server"
            raise ValueError(msg)

        arguments = {}
        for argument in self.prompt_arguments:
            if argument.required:
                if argument.name in kwargs:
                    arguments[argument.name] = kwargs[argument.name]
                else:
                    msg = f"Missing required argument {argument.name}"
                    raise RuntimeError(msg)
        return arguments

    @override
    def format_string(self, **arguments: str) -> str:
        """Get the prompt from the MCP server with optional arguments..

        Returns a text representation of the prompt.
        """
        text = "\n".join(
            message.content for message in self.format_messages(**arguments)
        )
        return text

    @override
    def format_messages(
        self, role: MessageRole = MessageRole.SYSTEM, **arguments: str
    ) -> List[ChatMessage]:
        """Get the prompt from the MCP server with optional arguments.

        Returns a list of messages formatted according to the MCP server.
        """
        arguments = self._check_arguments(**arguments)
        return self.mcp_server.get_prompt(self.name, arguments)

    @override
    def close(self) -> None:
        if self.mcp_server is not None:
            try:
                self.mcp_server.close()
            finally:
                self.mcp_server = None


class MCPServer(Resource, ABC):
    """Resource representing an MCP server and exposing its tools/prompts.

    This is a logical container for MCP tools and prompts; it is not directly invokable.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    # HTTP connection parameters
    endpoint: str
    headers: Dict[str, Any] | None = None
    timeout: int = 30
    sse_read_timeout: int = 60 * 5
    auth: httpx.Auth | None = None

    session: ClientSession = Field(default=None, exclude=True)
    connection_context: AsyncExitStack = Field(default=None, exclude=True)

    @field_serializer("auth")
    def __serialize_auth(self, auth: httpx.Auth | None) -> str | None:
        if auth is None:
            return auth
        return base64.b64encode(cloudpickle.dumps(auth)).decode()

    @model_validator(mode="before")
    def __custom_deserialize(self) -> "MCPServer":
        if "auth" in self:
            auth = self["auth"]
            if auth and isinstance(auth, str):
                self["auth"] = cloudpickle.loads(base64.b64decode(auth))
        return self

    @classmethod
    @override
    def resource_type(cls) -> ResourceType:
        return ResourceType.MCP_SERVER

    def _is_valid_http_url(self) -> bool:
        """Check if the endpoint is a valid HTTP URL."""
        try:
            parsed = urlparse(self.endpoint)
            return parsed.scheme in ("http", "https") and bool(parsed.netloc)
        except Exception:
            return False

    @asynccontextmanager
    async def _get_session(self) -> AsyncIterator[ClientSession]:
        """Get or create an MCP client session using the official SDK's streamable
        HTTP transport.
        """
        if not self._is_valid_http_url():
            msg = f"Invalid HTTP endpoint: {self.endpoint}"
            raise ValueError(msg)

        async with (
            streamablehttp_client(
                url=self.endpoint,
                headers=self.headers,
                timeout=timedelta(seconds=self.timeout),
                sse_read_timeout=timedelta(seconds=self.sse_read_timeout),
                auth=self.auth,
            ) as (read, write, _),
            ClientSession(read, write) as session,
        ):
            await session.initialize()
            yield session

    async def _cleanup_connection(self) -> None:
        """Clean up connection resources."""
        try:
            if self.connection_context is not None:
                await self.connection_context.aclose()
                self.connection_context = None
            self.session = None
        except Exception:
            pass

    async def call_tool_async(self, tool_name: str, *args: Any, **kwargs: Any) -> Any:
        """Call a tool on the MCP server asynchronously."""
        async with self._get_session() as session:
            arguments = kwargs if kwargs else (args[0] if args else {})

            result = await session.call_tool(
                tool_name,
                arguments,
                read_timeout_seconds=timedelta(seconds=self.timeout),
            )

            content = [extract_mcp_content_item(item) for item in result.content]

        return content

    def list_tools(self) -> List[MCPTool]:
        """List available tool definitions from the MCP server."""
        return asyncio.run(self._list_tools_async())

    async def _list_tools_async(self) -> List[MCPTool]:
        """Async implementation of list_tools."""
        async with self._get_session() as session:
            tools_response = await session.list_tools()

            tools = []
            for tool_data in tools_response.tools or []:
                metadata = ToolMetadata(
                    name=tool_data.name,
                    description=tool_data.description or "",
                    args_schema=tool_data.inputSchema
                    or {"type": "object", "properties": {}},
                )

                tool = MCPTool(
                    metadata=metadata,
                    mcp_server=self,
                )
                tools.append(tool)

            return tools

    def get_tool(self, name: str) -> MCPTool:
        """Get a single callable tool by name."""
        return asyncio.run(self._get_tool_async(name))

    async def _get_tool_async(self, name: str) -> MCPTool:
        """Async implementation of get_tool."""
        tools = await self._list_tools_async()
        for tool in tools:
            if tool.metadata.name == name:
                return tool
        msg = f"Tool '{name}' not found on MCP server at {self.endpoint}"
        raise ValueError(msg)

    def get_tool_metadata(self, name: str) -> ToolMetadata:
        """Get a single tool definition metadata by name."""
        tool = self.get_tool(name)
        return tool.metadata

    def list_prompts(self) -> List[MCPPrompt]:
        """List available prompts from the MCP server."""
        return asyncio.run(self._list_prompts_async())

    async def _list_prompts_async(self) -> List[MCPPrompt]:
        """Async implementation of list_prompts."""
        async with self._get_session() as session:
            # Check if the server supports prompts capability
            capabilities = session.get_server_capabilities()
            if capabilities is None or capabilities.prompts is None:
                # Server does not support prompts
                return []

            prompts_response = await session.list_prompts()

            prompts = []
            for prompt_data in prompts_response.prompts or []:
                arguments_dict = {}
                if prompt_data.arguments:
                    for arg in prompt_data.arguments:
                        arguments_dict[arg.name] = {
                            "description": arg.description,
                            "required": getattr(arg, "required", False),
                        }

                prompt = MCPPrompt(
                    name=prompt_data.name,
                    title=prompt_data.title,
                    description=prompt_data.description,
                    prompt_arguments=prompt_data.arguments,
                    mcp_server=self,
                )
                prompts.append(prompt)

            return prompts

    def get_prompt(
        self, name: str, arguments: Dict[str, str] | None = None
    ) -> List[ChatMessage]:
        """Get a single prompt definition by name."""
        return asyncio.run(self._get_prompt_async(name, arguments))

    async def _get_prompt_async(
        self, name: str, arguments: Dict[str, str] | None = None
    ) -> List[ChatMessage]:
        """Async implementation of get_prompt."""
        async with self._get_session() as session:
            prompt = await session.get_prompt(name, arguments)
            chat_messages = []
            for message in prompt.messages:
                if isinstance(message.content, TextContent):
                    chat_messages.append(
                        ChatMessage(role=message.role, content=message.content.text)
                    )
                else:
                    err_msg = f"Unsupported content type: {type(message.content)}"
                    raise TypeError(err_msg)

        return chat_messages

    @override
    def close(self) -> None:
        """Close the MCP server connection and clean up resources."""
        asyncio.run(self._cleanup_connection())
