#!/usr/bin/env python3
"""
MCP UI Server - provides interactive UI rendering tools.
Uses mcp-ui-server SDK to create standard UIResource objects.
"""

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,
)


ASK_USER_MARKER = "__ask_user__"


def ask_user(question: str, options: list = None) -> Dict[str, Any]:
    """Create an ask_user response.

    Returns a JSON structure with a marker so the backend can detect it
    and emit it as a special delta.ask_user event at the end of the stream.
    """
    payload = {
        "__type__": ASK_USER_MARKER,
        "question": question,
        "options": options or [],
    }
    return {
        "content": [
            {"type": "text", "text": json.dumps(payload, ensure_ascii=False)}
        ]
    }


def render_ui(
    title: str, html_content: str, width: str = "100%", height: str = "400px"
) -> Dict[str, Any]:
    """Create a UI resource and serialize it as JSON text for passthrough.

    The UIResource is serialized as a JSON string inside a type:"text" content block.
    This is necessary because langchain_mcp_adapters strips metadata (uri, mimeType)
    from EmbeddedResource objects. By wrapping it as text, the full resource JSON
    passes through to the frontend for detection and rendering.
    """
    try:
        uri_slug = title.replace(" ", "-").lower()[:50]
        ui_resource = create_ui_resource(
            {
                "uri": f"ui://mcp-ui-skill/{uri_slug}",
                "content": {"type": "rawHtml", "htmlString": html_content},
                "encoding": "text",
                "uiMetadata": {
                    UIMetadataKey.PREFERRED_FRAME_SIZE: [width, height],
                },
            }
        )

        # Serialize the full UIResource as JSON string for passthrough
        resource_json = json.dumps(
            ui_resource.model_dump(mode="json"), ensure_ascii=False
        )

        return {"content": [{"type": "text", "text": resource_json}]}

    except Exception as e:
        return {
            "content": [
                {"type": "text", "text": f"Error creating UI resource: {str(e)}"}
            ]
        }


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 HTML 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.",
                                },
                            },
                            "required": ["title", "html_content"],
                        },
                    }
                ]
            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":
                title = arguments.get("title", "UI Widget")
                html_content = arguments.get("html_content", "")
                width = arguments.get("width", "100%")
                height = arguments.get("height", "400px")

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

                result = render_ui(title, html_content, width, height)
                return {"jsonrpc": "2.0", "id": request_id, "result": result}

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

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

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