#!/usr/bin/env python3
"""
MCP UI Server - provides interactive UI rendering tools.
"""

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 ask_user() -> Dict[str, Any]:
    """Return a UIResource response for ask_user tool.

    The actual questions are in the TOOL_CALL arguments. The frontend
    detects ui_type from the UIResource metadata and extracts content
    from the corresponding TOOL_CALL args.
    """
    resource = create_ui_resource({
        "uri": "ui://mcp-ui/ask-user",
        "content": {"type": "rawHtml", "htmlString": "Questions sent to user."},
        "encoding": "text",
        "uiMetadata": {
            "type": "ask_user",
            "interactive": True,
        },
    })
    return {"content": [{"type": "text", "text": _serialize_ui_resource(resource)}]}


def render_ui(arguments: Dict[str, Any]) -> Dict[str, Any]:
    """Return a UIResource response for render_ui tool.

    The actual html_content/url is in the TOOL_CALL arguments. The frontend
    detects ui_type from the UIResource metadata and extracts content
    from the corresponding TOOL_CALL args.
    """
    html_content = arguments.get("html_content", "")
    url = arguments.get("url", "")
    width = arguments.get("width", "100%")
    height = arguments.get("height", "auto")

    if html_content:
        resource = create_ui_resource({
            "uri": "ui://mcp-ui/render-ui",
            "content": {"type": "rawHtml", "htmlString": "UI rendered."},
            "encoding": "text",
            "uiMetadata": {
                UIMetadataKey.PREFERRED_FRAME_SIZE: [width, height],
                "type": "render_ui_html",
                "interactive": False,
            },
        })
    else:
        resource = create_ui_resource({
            "uri": "ui://mcp-ui/render-ui",
            "content": {"type": "externalUrl", "iframeUrl": url},
            "encoding": "text",
            "uiMetadata": {
                UIMetadataKey.PREFERRED_FRAME_SIZE: [width, height],
                "type": "render_ui_url",
                "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 widget in the chat.",
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "title": {
                                    "type": "string",
                                    "description": "A descriptive title for the UI widget",
                                },
                                "html_content": {
                                    "type": "string",
                                    "description": "Complete HTML content to render. Use this OR url.",
                                },
                                "url": {
                                    "type": "string",
                                    "description": "External URL to embed. Use this OR html_content.",
                                },
                            },
                            "required": ["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":
                html_content = arguments.get("html_content", "")
                url = arguments.get("url", "")

                if not html_content and not url:
                    return create_error_response(
                        request_id, -32602, "Missing required parameter: html_content or url"
                    )

                result = render_ui(arguments)
                return {"jsonrpc": "2.0", "id": request_id, "result": result}

            elif tool_name == "ask_user":
                questions = arguments.get("questions", [])

                if not questions:
                    return create_error_response(
                        request_id, -32602, "Missing required parameter: questions"
                    )

                result = ask_user()
                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())