#!/usr/bin/env python3
"""
MCP UI Server - provides interactive UI rendering tools.
Uses URI-based routing: ui://mcp-ui/[resource-type]
"""

import asyncio
import json
import sys
from typing import Any, Dict

from mcp_ui_server import create_ui_resource, UIMetadataKey

from mcp_common import (
    create_error_response,
    create_initialize_response,
    create_ping_response,
    create_tools_list_response,
    load_tools_from_json,
    handle_mcp_streaming,
)


def _serialize_ui_resource(ui_resource) -> str:
    """Serialize a UIResource to JSON string."""
    return json.dumps(ui_resource.model_dump(mode="json"), ensure_ascii=False)


def _handle_render_ui(uri: str, data: Dict[str, Any]) -> Dict[str, Any]:
    """Unified handler for all render_ui URIs.

    Content is NOT self-contained — the actual html_content/url/questions
    stays in the TOOL_CALL args to save tokens. resource.text is a placeholder.
    The frontend extracts content from TOOL_CALL args based on URI.
    """
    width = data.get("width", "100%")
    height = data.get("height", "auto")

    if uri == "ui://mcp-ui/ask-user":
        resource = create_ui_resource({
            "uri": uri,
            "content": {"type": "rawHtml", "htmlString": "Questions sent to user."},
            "encoding": "text",
            "uiMetadata": {
                "interactive": True,
            },
        })
    elif uri == "ui://mcp-ui/url":
        url = data.get("url", "")
        resource = create_ui_resource({
            "uri": uri,
            "content": {"type": "externalUrl", "iframeUrl": url},
            "encoding": "text",
            "uiMetadata": {
                UIMetadataKey.PREFERRED_FRAME_SIZE: [width, height],
                "interactive": False,
            },
        })
    else:
        # ui://mcp-ui/html (default)
        resource = create_ui_resource({
            "uri": uri,
            "content": {"type": "rawHtml", "htmlString": "UI rendered."},
            "encoding": "text",
            "uiMetadata": {
                UIMetadataKey.PREFERRED_FRAME_SIZE: [width, height],
                "interactive": False,
            },
        })

    return {"content": [{"type": "text", "text": _serialize_ui_resource(resource)}]}


async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
    """Handle an MCP request."""
    try:
        method = request.get("method")
        params = request.get("params", {})
        request_id = request.get("id")

        if method == "initialize":
            return create_initialize_response(request_id, "mcp-ui")

        elif method == "ping":
            return create_ping_response(request_id)

        elif method == "tools/list":
            tools = load_tools_from_json("mcp_ui_tools.json")
            if not tools:
                tools = [
                    {
                        "name": "render_ui",
                        "description": "Render an interactive UI resource in the chat.",
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "uri": {"type": "string"},
                                "title": {"type": "string"},
                                "html_content": {"type": "string"},
                                "url": {"type": "string"},
                            },
                            "required": ["uri", "title"],
                        },
                    }
                ]
            return create_tools_list_response(request_id, tools)

        elif method == "tools/call":
            tool_name = params.get("name")
            arguments = params.get("arguments", {})

            if tool_name == "render_ui":
                uri = arguments.get("uri", "")
                data = arguments.get("data", {})

                if not uri or not uri.startswith("ui://mcp-ui/"):
                    return create_error_response(
                        request_id, -32602,
                        "Invalid uri. Must be one of: ui://mcp-ui/html, ui://mcp-ui/url, ui://mcp-ui/ask-user"
                    )

                path = uri.replace("ui://mcp-ui/", "")

                if path == "html" and not data.get("html_content"):
                    return create_error_response(
                        request_id, -32602, "Missing required parameter: data.html_content (for uri=ui://mcp-ui/html)"
                    )
                elif path == "url" and not data.get("url"):
                    return create_error_response(
                        request_id, -32602, "Missing required parameter: data.url (for uri=ui://mcp-ui/url)"
                    )
                elif path == "ask-user" and not data.get("questions"):
                    return create_error_response(
                        request_id, -32602, "Missing required parameter: data.questions (for uri=ui://mcp-ui/ask-user)"
                    )

                result = _handle_render_ui(uri, data)
                return {"jsonrpc": "2.0", "id": request_id, "result": result}

            else:
                return create_error_response(
                    request_id, -32601, f"Unknown tool: {tool_name}"
                )

        else:
            return create_error_response(
                request_id, -32601, f"Unknown method: {method}"
            )

    except Exception as e:
        return create_error_response(
            request.get("id"), -32603, f"Internal error: {str(e)}"
        )


async def main():
    """Main entry point."""
    await handle_mcp_streaming(handle_request)


if __name__ == "__main__":
    asyncio.run(main())