"""MCP server exposing Skytale encrypted channels.

Run as::

    SKYTALE_IDENTITY=my-agent python -m skytale_sdk.integrations._mcp

Works with any MCP client: Claude Desktop, LangGraph, CrewAI, etc.

Transport options:
    - ``stdio`` (default): for local MCP clients
    - ``streamable-http``: Streamable HTTP endpoint (replaces deprecated SSE)

Authentication:
    - API key via ``SKYTALE_API_KEY`` env var (for Skytale relay)
    - OAuth 2.1 with PKCE for MCP Streamable HTTP (configured via env vars
      or :class:`SkytaleToolServer` constructor)
"""

from __future__ import annotations

import os
from typing import Dict, List, Optional

from mcp.server.fastmcp import FastMCP

from skytale_sdk.channels import SkytaleChannelManager

_identity = os.environ.get("SKYTALE_IDENTITY", "mcp-agent")
_manager = SkytaleChannelManager(identity=_identity.encode())

mcp = FastMCP("skytale")


@mcp.tool()
def skytale_create_channel(channel: str) -> str:
    """Create a new encrypted channel.

    Creates an MLS-encrypted channel and starts listening for messages.

    Args:
        channel: Channel name in org/namespace/service format
            (e.g. "acme/research/results").
    """
    _manager.create(channel)
    return f"Channel {channel} created"


@mcp.tool()
def skytale_key_package() -> str:
    """Generate an MLS key package for joining a channel.

    Returns a hex-encoded key package. Send this to the channel creator
    so they can call skytale_add_member to add you.
    """
    kp = _manager.key_package()
    return kp.hex()


@mcp.tool()
def skytale_add_member(channel: str, key_package_hex: str) -> str:
    """Add a member to a channel using their key package.

    Args:
        channel: Channel name to add the member to.
        key_package_hex: Hex-encoded MLS key package from skytale_key_package.

    Returns:
        Hex-encoded MLS Welcome message. Send this to the joining agent
        so they can call skytale_join_channel.
    """
    kp = bytes.fromhex(key_package_hex)
    welcome = _manager.add_member(channel, kp)
    return welcome.hex()


@mcp.tool()
def skytale_join_channel(channel: str, welcome_hex: str) -> str:
    """Join an existing channel using an MLS Welcome message.

    Args:
        channel: Channel name to join.
        welcome_hex: Hex-encoded Welcome message from skytale_add_member.
    """
    welcome = bytes.fromhex(welcome_hex)
    _manager.join(channel, welcome)
    return f"Joined channel {channel}"


@mcp.tool()
def skytale_send(channel: str, message: str) -> str:
    """Send an encrypted message to a channel.

    Messages are end-to-end encrypted using the MLS protocol. All members
    of the channel will receive the message.

    Args:
        channel: Channel name in org/namespace/service format.
        message: Message text to send.
    """
    _manager.send(channel, message)
    return f"Message sent to {channel}"


@mcp.tool()
def skytale_receive(channel: str, timeout: float = 5.0) -> str:
    """Receive messages from an encrypted channel.

    Returns all messages received since the last check, or waits up to
    timeout seconds for new messages.

    Args:
        channel: Channel name in org/namespace/service format.
        timeout: Seconds to wait for messages (default 5).
    """
    msgs = _manager.receive(channel, timeout=timeout)
    if not msgs:
        return f"No new messages on {channel}"
    return "\n".join(f"[{channel}] {m}" for m in msgs)


@mcp.tool()
def skytale_channels() -> str:
    """List all active encrypted channels."""
    channels = _manager.list_channels()
    if not channels:
        return "No active channels"
    return ", ".join(channels)


class SkytaleToolServer:
    """Configurable MCP server wrapping Skytale encrypted channel tools.

    Use this when you need programmatic control over identity, transport,
    or OAuth configuration instead of the module-level ``mcp`` instance.

    Supports two transports:
    - ``stdio``: for local MCP clients (default)
    - ``streamable-http``: Streamable HTTP endpoint with optional OAuth 2.1

    Args:
        identity: Agent identity string.
        transport: Transport type: ``"stdio"`` or ``"streamable-http"``.
        host: Bind address for Streamable HTTP (default ``"127.0.0.1"``).
        port: Port for Streamable HTTP (default ``8080``).
        oauth_client_id: OAuth 2.1 client ID (for Streamable HTTP auth).
        oauth_authorization_url: OAuth authorization endpoint.
        oauth_token_url: OAuth token endpoint.
        oauth_scopes: OAuth scopes (comma-separated string or list).

    Example::

        server = SkytaleToolServer(
            identity="my-agent",
            transport="streamable-http",
            port=9090,
        )
        server.run()

        # With OAuth 2.1:
        server = SkytaleToolServer(
            identity="my-agent",
            transport="streamable-http",
            oauth_client_id="skytale-agent",
            oauth_authorization_url="https://auth.example.com/authorize",
            oauth_token_url="https://auth.example.com/token",
        )
        server.run()
    """

    def __init__(
        self,
        identity: str = "mcp-agent",
        transport: str = "stdio",
        host: str = "127.0.0.1",
        port: int = 8080,
        oauth_client_id: Optional[str] = None,
        oauth_authorization_url: Optional[str] = None,
        oauth_token_url: Optional[str] = None,
        oauth_scopes: Optional[List[str]] = None,
    ) -> None:
        self._identity = identity
        self._transport = transport
        self._host = host
        self._port = port
        self._oauth_client_id = oauth_client_id or os.environ.get(
            "MCP_OAUTH_CLIENT_ID"
        )
        self._oauth_authorization_url = oauth_authorization_url or os.environ.get(
            "MCP_OAUTH_AUTHORIZATION_URL"
        )
        self._oauth_token_url = oauth_token_url or os.environ.get(
            "MCP_OAUTH_TOKEN_URL"
        )
        if oauth_scopes is not None:
            self._oauth_scopes = oauth_scopes
        else:
            env_scopes = os.environ.get("MCP_OAUTH_SCOPES", "")
            self._oauth_scopes = (
                [s.strip() for s in env_scopes.split(",") if s.strip()]
                if env_scopes
                else ["mcp"]
            )

        self._manager = SkytaleChannelManager(identity=identity.encode())
        self._server = FastMCP("skytale")
        self._register_tools()

    @property
    def oauth_configured(self) -> bool:
        """Return True if all OAuth 2.1 parameters are set.

        Example::

            >>> server = SkytaleToolServer()
            >>> server.oauth_configured
            False
        """
        return all([
            self._oauth_client_id,
            self._oauth_authorization_url,
            self._oauth_token_url,
        ])

    def _register_tools(self) -> None:
        """Register Skytale channel tools on the FastMCP server instance."""
        mgr = self._manager

        @self._server.tool()
        def skytale_create_channel(channel: str) -> str:
            """Create a new encrypted channel.

            Creates an MLS-encrypted channel and starts listening for messages.

            Args:
                channel: Channel name in org/namespace/service format.
            """
            mgr.create(channel)
            return f"Channel {channel} created"

        @self._server.tool()
        def skytale_key_package() -> str:
            """Generate an MLS key package for joining a channel.

            Returns a hex-encoded key package.
            """
            kp = mgr.key_package()
            return kp.hex()

        @self._server.tool()
        def skytale_add_member(channel: str, key_package_hex: str) -> str:
            """Add a member to a channel using their key package.

            Args:
                channel: Channel name to add the member to.
                key_package_hex: Hex-encoded MLS key package.

            Returns:
                Hex-encoded MLS Welcome message.
            """
            kp = bytes.fromhex(key_package_hex)
            welcome = mgr.add_member(channel, kp)
            return welcome.hex()

        @self._server.tool()
        def skytale_join_channel(channel: str, welcome_hex: str) -> str:
            """Join an existing channel using an MLS Welcome message.

            Args:
                channel: Channel name to join.
                welcome_hex: Hex-encoded Welcome message.
            """
            welcome = bytes.fromhex(welcome_hex)
            mgr.join(channel, welcome)
            return f"Joined channel {channel}"

        @self._server.tool()
        def skytale_send(channel: str, message: str) -> str:
            """Send an encrypted message to a channel.

            Args:
                channel: Channel name in org/namespace/service format.
                message: Message text to send.
            """
            mgr.send(channel, message)
            return f"Message sent to {channel}"

        @self._server.tool()
        def skytale_receive(channel: str, timeout: float = 5.0) -> str:
            """Receive messages from an encrypted channel.

            Args:
                channel: Channel name in org/namespace/service format.
                timeout: Seconds to wait for messages (default 5).
            """
            msgs = mgr.receive(channel, timeout=timeout)
            if not msgs:
                return f"No new messages on {channel}"
            return "\n".join(f"[{channel}] {m}" for m in msgs)

        @self._server.tool()
        def skytale_channels() -> str:
            """List all active encrypted channels."""
            channels = mgr.list_channels()
            if not channels:
                return "No active channels"
            return ", ".join(channels)

    def run(self) -> None:
        """Start the MCP server with the configured transport.

        For ``stdio`` transport, blocks reading from stdin.
        For ``streamable-http`` transport, starts an HTTP server on the
        configured host and port.

        Example::

            server = SkytaleToolServer(transport="stdio")
            server.run()
        """
        if self._transport == "streamable-http":
            self._server.run(
                transport="streamable-http",
                host=self._host,
                port=self._port,
            )
        else:
            self._server.run(transport="stdio")


class SkytaleToolClient:
    """Helper for connecting to a Skytale MCP server via Streamable HTTP.

    Wraps the connection setup for MCP clients that want to talk to a
    :class:`SkytaleToolServer` running in ``streamable-http`` mode, with
    optional OAuth 2.1 authentication and MLS encryption of the transport.

    Args:
        url: MCP server Streamable HTTP endpoint URL.
        channel_manager: Optional :class:`SkytaleChannelManager` for MLS
            encryption of the transport layer.  When provided, all JSON-RPC
            messages are encrypted before transmission.
        channel: Channel name for encrypted transport (required if
            ``channel_manager`` is provided).
        oauth_client_id: OAuth 2.1 client ID.
        oauth_authorization_url: OAuth authorization endpoint.
        oauth_token_url: OAuth token endpoint.
        oauth_scopes: OAuth scopes.

    Example::

        client = SkytaleToolClient(
            url="https://mcp.example.com/mcp",
        )
        transport = client.create_transport()

        # With encryption:
        mgr = SkytaleChannelManager(identity=b"caller")
        mgr.create("org/ns/mcp-rpc")
        client = SkytaleToolClient(
            url="https://mcp.example.com/mcp",
            channel_manager=mgr,
            channel="org/ns/mcp-rpc",
        )
        transport = client.create_transport()
    """

    def __init__(
        self,
        url: str,
        channel_manager: Optional[SkytaleChannelManager] = None,
        channel: Optional[str] = None,
        oauth_client_id: Optional[str] = None,
        oauth_authorization_url: Optional[str] = None,
        oauth_token_url: Optional[str] = None,
        oauth_scopes: Optional[List[str]] = None,
    ) -> None:
        self._url = url
        self._manager = channel_manager
        self._channel = channel
        self._oauth_client_id = oauth_client_id
        self._oauth_authorization_url = oauth_authorization_url
        self._oauth_token_url = oauth_token_url
        self._oauth_scopes = oauth_scopes or ["mcp"]

    @property
    def oauth_config(self):
        """Build an :class:`OAuthConfig` from the client's OAuth parameters.

        Returns:
            An :class:`OAuthConfig` if all required fields are set, else ``None``.

        Example::

            cfg = client.oauth_config
            if cfg:
                url, verifier = build_authorization_url(cfg)
        """
        if not all([
            self._oauth_client_id,
            self._oauth_authorization_url,
            self._oauth_token_url,
        ]):
            return None

        from skytale_sdk.integrations._mcp_transport import OAuthConfig

        return OAuthConfig(
            client_id=self._oauth_client_id,
            authorization_url=self._oauth_authorization_url,
            token_url=self._oauth_token_url,
            scopes=self._oauth_scopes,
        )

    def create_transport(self):
        """Create an :class:`EncryptedStreamableHTTPTransport` for this client.

        If a channel manager and channel are configured, the transport uses
        MLS encryption.  If OAuth parameters are set, the transport handles
        token management automatically.

        Returns:
            An :class:`EncryptedStreamableHTTPTransport` instance.

        Example::

            transport = client.create_transport()
            async with transport.connect() as (read, write):
                await write({"jsonrpc": "2.0", "method": "ping", "id": 1})
                resp = await read()
        """
        from skytale_sdk.integrations._mcp_transport import (
            EncryptedStreamableHTTPTransport,
        )

        return EncryptedStreamableHTTPTransport(
            url=self._url,
            channel_manager=self._manager or SkytaleChannelManager(
                identity=b"mcp-client",
                mock=True,
            ),
            channel=self._channel or "__mcp_default__",
            oauth=self.oauth_config,
        )


if __name__ == "__main__":
    import sys

    transport = "stdio"
    if "--streamable-http" in sys.argv:
        transport = "streamable-http"

    if transport == "streamable-http":
        server = SkytaleToolServer(
            identity=_identity,
            transport="streamable-http",
        )
        server.run()
    else:
        mcp.run(transport="stdio")
