"""FastMCP conformance test server.

Registers the exact tools, resources, and prompts expected by the
MCP conformance test suite (https://github.com/modelcontextprotocol/conformance).
"""

import asyncio
import base64
import json
import sys
from enum import Enum as PyEnum

import mcp_types
from mcp_types import EmbeddedResource, ImageContent, TextContent
from pydantic import BaseModel, Field

from fastmcp import FastMCP
from fastmcp.exceptions import ToolError
from fastmcp.prompts import Message
from fastmcp.server.context import Context
from fastmcp.tools.function_tool import FunctionTool
from fastmcp.utilities.types import Audio, Image

# Minimal 1x1 red PNG for image tests (89 bytes)
_1X1_PNG = base64.b64decode(
    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4"
    "nGP4z8BQDwAEgAF/pooBPQAAAABJRU5ErkJggg=="
)

# Minimal valid WAV: 16-bit mono PCM, 44100 Hz, single silent sample
_SILENT_WAV = (
    b"RIFF"
    + (38).to_bytes(4, "little")
    + b"WAVEfmt "
    + (16).to_bytes(4, "little")
    + (1).to_bytes(2, "little")  # PCM
    + (1).to_bytes(2, "little")  # mono
    + (44100).to_bytes(4, "little")  # sample rate
    + (88200).to_bytes(4, "little")  # byte rate
    + (2).to_bytes(2, "little")  # block align
    + (16).to_bytes(2, "little")  # bits per sample
    + b"data"
    + (2).to_bytes(4, "little")
    + (0).to_bytes(2, "little")  # one silent sample
)

server = FastMCP("conformance-test-server", dereference_schemas=False)


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------


@server.tool(name="test_simple_text")
async def test_simple_text() -> str:
    """A simple text tool for conformance testing."""
    return "This is a simple text response for testing."


@server.tool(name="test_image_content")
async def test_image_content() -> Image:
    """Returns a PNG image."""
    return Image(data=_1X1_PNG, format="png")


@server.tool(name="test_audio_content")
async def test_audio_content() -> Audio:
    """Returns WAV audio."""
    return Audio(data=_SILENT_WAV, format="wav")


@server.tool(name="test_embedded_resource")
async def test_embedded_resource() -> list:
    """Returns an embedded resource."""
    return [
        EmbeddedResource(
            type="resource",
            resource=mcp_types.TextResourceContents(
                uri="test://embedded-resource",
                mime_type="text/plain",
                text="This is an embedded resource content.",
            ),
        )
    ]


@server.tool(name="test_multiple_content_types")
async def test_multiple_content_types() -> list:
    """Returns mixed text, image, and resource content."""
    return [
        TextContent(type="text", text="This is a text part of the response."),
        ImageContent(
            type="image",
            data=base64.b64encode(_1X1_PNG).decode(),
            mime_type="image/png",
        ),
        EmbeddedResource(
            type="resource",
            resource=mcp_types.TextResourceContents(
                uri="test://mixed-content-resource",
                mime_type="application/json",
                text='{"test":"data","value":123}',
            ),
        ),
    ]


@server.tool(name="test_error_handling")
async def test_error_handling() -> str:
    """Always returns an error."""
    raise ToolError("This tool intentionally returns an error for testing")


@server.tool(name="test_tool_with_logging")
async def test_tool_with_logging(ctx: Context) -> str:
    """Sends log notifications during execution."""
    await ctx.info("Tool execution started")
    await asyncio.sleep(0.05)
    await ctx.info("Tool processing data")
    await asyncio.sleep(0.05)
    await ctx.info("Tool execution completed")
    return "Logging test complete."


@server.tool(name="test_tool_with_progress")
async def test_tool_with_progress(ctx: Context) -> str:
    """Reports progress notifications."""
    await ctx.report_progress(0, 100)
    await asyncio.sleep(0.05)
    await ctx.report_progress(50, 100)
    await asyncio.sleep(0.05)
    await ctx.report_progress(100, 100)
    return "Progress test complete."


@server.tool(name="test_sampling")
async def test_sampling(prompt: str, ctx: Context) -> str:
    """Requests LLM sampling via the client."""
    result = await ctx.sample(
        messages=[prompt],
        result_type=str,
    )
    return f"Sampling result: {result}"


class _UserInfo(BaseModel):
    username: str
    email: str


@server.tool(name="test_elicitation")
async def test_elicitation(message: str, ctx: Context) -> str:
    """Requests user input via elicitation."""
    result = await ctx.elicit(message, _UserInfo)
    return f"Elicitation result: {result}"


class _UserStatus(str, PyEnum):
    active = "active"
    inactive = "inactive"
    pending = "pending"


class _DefaultsForm(BaseModel):
    name: str = Field(default="John Doe", description="User name")
    age: int = Field(default=30, description="User age")
    score: float = Field(default=95.5, description="User score")
    status: _UserStatus = Field(default=_UserStatus.active, description="User status")
    verified: bool = Field(default=True, description="Verification status")


@server.tool(name="test_elicitation_sep1034_defaults")
async def test_elicitation_sep1034_defaults(ctx: Context) -> str:
    """Tests elicitation with default values per SEP-1034."""
    result = await ctx.elicit(
        "Please review and update the form fields with defaults",
        _DefaultsForm,
    )
    return f"Elicitation completed: {result}"


@server.tool(name="test_elicitation_sep1330_enums")
async def test_elicitation_sep1330_enums(ctx: Context) -> str:
    """Tests elicitation with enum schema improvements per SEP-1330."""
    result = await ctx.session.elicit(
        message="Please select options from the enum fields",
        requested_schema={
            "type": "object",
            "properties": {
                "untitledSingle": {
                    "type": "string",
                    "description": "Select one option",
                    "enum": ["option1", "option2", "option3"],
                },
                "titledSingle": {
                    "type": "string",
                    "description": "Select one option with titles",
                    "oneOf": [
                        {"const": "value1", "title": "First Option"},
                        {"const": "value2", "title": "Second Option"},
                        {"const": "value3", "title": "Third Option"},
                    ],
                },
                "legacyEnum": {
                    "type": "string",
                    "description": "Select one option (legacy)",
                    "enum": ["opt1", "opt2", "opt3"],
                    "enumNames": [
                        "Option One",
                        "Option Two",
                        "Option Three",
                    ],
                },
                "untitledMulti": {
                    "type": "array",
                    "description": "Select multiple options",
                    "minItems": 1,
                    "maxItems": 3,
                    "items": {
                        "type": "string",
                        "enum": ["option1", "option2", "option3"],
                    },
                },
                "titledMulti": {
                    "type": "array",
                    "description": "Select multiple options with titles",
                    "minItems": 1,
                    "maxItems": 3,
                    "items": {
                        "anyOf": [
                            {"const": "value1", "title": "First Choice"},
                            {"const": "value2", "title": "Second Choice"},
                            {"const": "value3", "title": "Third Choice"},
                        ]
                    },
                },
            },
            "required": [],
        },
        related_request_id=ctx.request_id,
    )
    return f"Elicitation completed: action={result.action}, content={json.dumps(result.content or {})}"


async def _json_schema_2020_12_fn(
    name: str | None = None,
    address: dict | None = None,
) -> str:
    """Tool with JSON Schema 2020-12 features for conformance testing (SEP-1613)."""
    return f"JSON Schema 2020-12 tool called with: name={name}, address={address}"


server.add_tool(
    FunctionTool(
        fn=_json_schema_2020_12_fn,
        name="json_schema_2020_12_tool",
        description="Tool with JSON Schema 2020-12 features for conformance testing (SEP-1613)",
        parameters={
            "$schema": "https://json-schema.org/draft/2020-12/schema",
            "type": "object",
            "$defs": {
                "address": {
                    "type": "object",
                    "properties": {
                        "street": {"type": "string"},
                        "city": {"type": "string"},
                    },
                }
            },
            "properties": {
                "name": {"type": "string"},
                "address": {"$ref": "#/$defs/address"},
            },
            "additionalProperties": False,
        },
    )
)


# ---------------------------------------------------------------------------
# Resources
# ---------------------------------------------------------------------------


@server.resource(
    "test://static-text",
    name="Static text resource",
    mime_type="text/plain",
)
async def static_text_resource() -> str:
    """Returns static text content."""
    return "This is the content of the static text resource."


@server.resource(
    "test://static-binary",
    name="Static binary resource",
    mime_type="image/png",
)
async def static_binary_resource() -> bytes:
    """Returns a binary PNG image."""
    return _1X1_PNG


@server.resource(
    "test://template/{id}/data",
    name="Template resource",
    mime_type="application/json",
)
async def template_resource(id: str) -> str:
    """Returns JSON data with the template parameter substituted."""
    return json.dumps({"id": id, "templateTest": True, "data": f"Data for ID: {id}"})


@server.resource(
    "test://watched-resource",
    name="Watched resource",
    mime_type="text/plain",
)
async def watched_resource() -> str:
    """A resource that supports subscriptions."""
    return "Watched resource content."


# ---------------------------------------------------------------------------
# Prompts
# ---------------------------------------------------------------------------


@server.prompt(name="test_simple_prompt")
async def test_simple_prompt() -> str:
    """A simple prompt for conformance testing."""
    return "This is a simple prompt for testing."


@server.prompt(name="test_prompt_with_arguments")
async def test_prompt_with_arguments(arg1: str, arg2: str) -> str:
    """A prompt that accepts arguments."""
    return f"Prompt with arguments: arg1='{arg1}', arg2='{arg2}'"


@server.prompt(name="test_prompt_with_embedded_resource")
async def test_prompt_with_embedded_resource(resourceUri: str) -> list:
    """A prompt that returns an embedded resource."""
    return [
        Message(
            EmbeddedResource(
                type="resource",
                resource=mcp_types.TextResourceContents(
                    uri=resourceUri,
                    mime_type="text/plain",
                    text=f"Content of resource {resourceUri}",
                ),
            )
        ),
    ]


@server.prompt(name="test_prompt_with_image")
async def test_prompt_with_image() -> list:
    """A prompt that returns an image."""
    return [
        Message(
            ImageContent(
                type="image",
                data=base64.b64encode(_1X1_PNG).decode(),
                mime_type="image/png",
            )
        ),
        Message("Please analyze the image above."),
    ]


if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8000
    server.run(transport="streamable-http", host="127.0.0.1", port=port)