"""Typed structured output, matching the Archipelago MCP server convention.

Every tool returns the exact text an agent reads, and carries alongside it a
Pydantic model holding the same facts. The server publishes the model as the
tool's ``outputSchema`` and returns it as ``structuredContent``.

This mirrors ``mcp_servers/filesystem/.../utils/structured_output.py`` in the
Archipelago repo deliberately: an environment that follows the house pattern
can be reviewed against the servers already in the harness instead of being
read from scratch.

Payloads are dumped with ``exclude_none`` because Archipelago flattens served
schemas, turning ``X | None`` into a bare ``X`` plus a non-standard
``nullable`` flag - an emitted ``null`` would then fail client-side
validation, while an omitted key always validates.
"""

from __future__ import annotations

import functools
import inspect
from collections.abc import Awaitable, Callable
from typing import Any

from fastmcp.tools import ToolResult
from mcp.types import TextContent
from pydantic import BaseModel


class StructuredText[M: BaseModel](str):
    """The exact text a tool returns, carrying its structured equivalent."""

    data: M

    def __new__(cls, text: str, data: M) -> "StructuredText[M]":
        obj = super().__new__(cls, text)
        obj.data = data
        return obj


def output_schema(model: type[BaseModel]) -> dict[str, Any]:
    return model.model_json_schema()


def to_tool_result(result: StructuredText[Any]) -> ToolResult:
    if not isinstance(result, StructuredText):
        raise TypeError(
            f"structured tool returned {type(result).__name__}, expected StructuredText"
        )
    return ToolResult(
        content=[TextContent(type="text", text=str(result))],
        structured_content=result.data.model_dump(mode="json", exclude_none=True),
    )


def structured_tool[**P](
    fn: Callable[P, Awaitable[StructuredText[Any]]],
) -> Callable[P, Awaitable[ToolResult]]:
    """Wrap a tool so the server returns its text plus structured content."""

    @functools.wraps(fn)
    async def tool(*args: P.args, **kwargs: P.kwargs) -> ToolResult:
        return to_tool_result(await fn(*args, **kwargs))

    setattr(  # noqa: B010
        tool, "__signature__",
        inspect.signature(fn).replace(return_annotation=ToolResult),
    )
    return tool