# Copyright (c) 2026 Dedalus Labs, Inc. and its contributors
# SPDX-License-Identifier: MIT

"""JSON Schema generation utilities for MCP tool arguments."""

from __future__ import annotations

from collections.abc import Iterable, Mapping, MutableMapping
from dataclasses import dataclass
from typing import Any

from pydantic import TypeAdapter
from pydantic.json_schema import JsonSchemaMode, JsonSchemaValue


__all__ = [
    "JsonSchema",
    "SchemaError",
    "SchemaEnvelope",
    "DEDALUS_BOX_KEY",
    "DEFAULT_WRAP_FIELD",
    "compress_schema",
    "enforce_strict_schema",
    "generate_schema_from_annotation",
    "ensure_object_schema",
    "unwrap_structured_content",
    "resolve_input_schema",
    "resolve_output_schema",
]


JsonSchema = JsonSchemaValue


_MISSING = object()

DEDALUS_BOX_KEY = "x-dedalus-box"
"""Vendor extension that records scalar boxing metadata."""


DEFAULT_WRAP_FIELD = "result"
"""Name of the object property used when auto-wrapping scalar schemas."""


class SchemaError(RuntimeError):
    """Raised when a schema cannot be generated or normalized."""


@dataclass(frozen=True, slots=True)
class SchemaEnvelope:
    """Describe how a JSON Schema aligns with MCP structured content.

    Attributes:
        schema: JSON Schema value, treated as immutable data.
        wrap_field: Name of the boxed scalar field, or ``None`` when the schema
            already describes an object.

    """

    schema: JsonSchema
    wrap_field: str | None = None

    @property
    def is_wrapped(self) -> bool:
        """Return whether the envelope records a boxed scalar.

        Returns:
            bool: ``True`` when :attr:`wrap_field` is not ``None``.

        """
        return self.wrap_field is not None

    def unwrap(self, structured_content: Mapping[str, Any]) -> Any:
        """Recover the original value from structured content.

        Args:
            structured_content: Mapping returned in
                :attr:`CallToolResult.structuredContent`.

        Returns:
            Any: Original scalar or mapping provided by the tool.

        Raises:
            SchemaError: If the structured content conflicts with the recorded
                boxing metadata.

        """
        if not self.is_wrapped:
            return structured_content

        if not isinstance(structured_content, Mapping):
            raise SchemaError("Structured content must be a mapping when using an auto-wrapped schema.")

        try:
            return structured_content[self.wrap_field]  # type: ignore[index]
        except KeyError as exc:
            raise SchemaError(f"Expected wrapped result to contain '{self.wrap_field}'") from exc

    def wrap(self, value: Any) -> Mapping[str, Any]:
        """Box ``value`` into the MCP structured-content envelope.

        Args:
            value: Original scalar or mapping.

        Returns:
            Mapping[str, Any]: Transport-ready payload.

        Raises:
            SchemaError: If the envelope records no boxing yet ``value`` is not
                a mapping.

        """
        if not self.is_wrapped:
            if isinstance(value, Mapping):
                return value
            raise SchemaError("Wrapped value must be a mapping when no synthetic wrapper is present.")
        return {self.wrap_field: value}


def generate_schema_from_annotation(
    annotation: Any,
    *,
    mode: JsonSchemaMode = "serialization",
    wrap_scalar: bool = True,
    wrap_field: str = DEFAULT_WRAP_FIELD,
    compress: bool = True,
    drop_titles: bool = True,
    relax_additional_properties: bool = True,
) -> SchemaEnvelope:
    """Build a schema envelope from a Python annotation.

    Args:
        annotation: Object understood by :class:`pydantic.TypeAdapter`.
        mode: JSON Schema generation mode.
        wrap_scalar: Whether non-object schemas should be boxed automatically.
        wrap_field: Synthetic property name used when boxing occurs.
        compress: Whether to remove cosmetic metadata.
        drop_titles: Remove ``title`` fields when compression is enabled.
        relax_additional_properties: Replace ``additionalProperties: false``
            with a permissive default when compression is enabled.

    Returns:
        SchemaEnvelope: Schema information aligned with MCP transport rules.

    Raises:
        SchemaError: If :class:`pydantic.TypeAdapter` cannot derive a schema.

    """
    return _coerce_envelope(
        annotation,
        wrap_scalar=wrap_scalar,
        wrap_field=wrap_field,
        mode=mode,
        compress=compress,
        drop_titles=drop_titles,
        relax_additional_properties=relax_additional_properties,
    )


def ensure_object_schema(
    schema: JsonSchema, *, wrap_scalar: bool = True, wrap_field: str = DEFAULT_WRAP_FIELD, marker: str = DEDALUS_BOX_KEY
) -> SchemaEnvelope:
    """Guarantee that ``schema`` can travel over MCP as an object.

    Args:
        schema: JSON Schema to inspect.
        wrap_scalar: Whether to box non-object schemas.
        wrap_field: Property name used when boxing.
        marker: Vendor extension recording boxing metadata.

    Returns:
        SchemaEnvelope: Schema aligned with MCP output rules.

    Raises:
        SchemaError: If boxing is disabled and ``schema`` is non-object.

    """
    if _describes_object(schema):
        return SchemaEnvelope(schema=_clone_schema(schema))

    if not wrap_scalar:
        raise SchemaError("Schema describes a non-object value. Set wrap_scalar=True to comply with MCP output rules.")

    wrapped: JsonSchema = {
        "type": "object",
        "properties": {wrap_field: _clone_schema(schema)},
        "required": [wrap_field],
        "additionalProperties": False,
        marker: {"field": wrap_field},
    }
    return SchemaEnvelope(schema=wrapped, wrap_field=wrap_field)


def resolve_input_schema(schema_like: Any) -> JsonSchema:
    """Normalize user-provided input schema declarations.

    Args:
        schema_like: JSON Schema mapping, :class:`SchemaEnvelope`, or Python type.

    Returns:
        JsonSchema: Object schema suitable for MCP ``inputSchema``.

    Raises:
        SchemaError: If the schema cannot be interpreted as an object schema.

    """
    envelope = _coerce_envelope(
        schema_like, wrap_scalar=False, wrap_field=DEFAULT_WRAP_FIELD, mode="validation", compress=True
    )
    return envelope.schema


def resolve_output_schema(schema_like: Any) -> SchemaEnvelope:
    """Normalize user-provided output schema declarations.

    Args:
        schema_like: JSON Schema mapping, :class:`SchemaEnvelope`, or Python type.

    Returns:
        SchemaEnvelope: Schema annotated with boxing metadata for MCP outputs.

    Raises:
        SchemaError: If the schema cannot be derived.

    """
    return _coerce_envelope(
        schema_like, wrap_scalar=True, wrap_field=DEFAULT_WRAP_FIELD, mode="serialization", compress=True
    )


def unwrap_structured_content(
    structured_content: Mapping[str, Any] | None,
    schema: Mapping[str, Any] | SchemaEnvelope,
    *,
    marker: str = DEDALUS_BOX_KEY,
) -> Any:
    """Reverse the boxing step recorded in a schema envelope.

    Args:
        structured_content: Payload returned by the remote tool.
        schema: Raw JSON Schema or :class:`SchemaEnvelope` describing boxing.
        marker: Vendor extension key signaling boxing.

    Returns:
        Any: Unboxed value.

    """
    if structured_content is None:
        return None

    envelope = schema if isinstance(schema, SchemaEnvelope) else _envelope_from_schema(schema, marker)
    return envelope.unwrap(structured_content)


def compress_schema(
    schema: JsonSchema,
    *,
    drop_titles: bool = True,
    relax_additional_properties: bool = True,
    prune_parameters: Iterable[str] | None = None,
) -> JsonSchema:
    """Return a structurally equivalent schema with cosmetic noise removed.

    Args:
        schema: JSON Schema to normalize.
        drop_titles: Remove ``title`` keys recursively.
        relax_additional_properties: Drop ``additionalProperties: false``.
        prune_parameters: Parameter names to delete from the top level.

    Returns:
        JsonSchema: Cleaned schema.

    """
    clone = _clone_schema(schema)

    if drop_titles:
        _strip_field(clone, "title")

    if relax_additional_properties:
        _relax_additional_properties(clone)

    if prune_parameters:
        for param in prune_parameters:
            _drop_top_level_property(clone, param)

    _prune_empty_required(clone)
    return clone


def enforce_strict_schema(schema: JsonSchema) -> JsonSchema:
    """Return a schema compliant with strict MCP/LLM expectations.

    The transformation ensures object schemas explicitly forbid unknown
    properties, populates ``required`` lists, cleans up oneOf/anyOf usage, and
    strips ``None`` defaults that add no semantic value.

    Args:
        schema: JSON Schema to normalize strictly.

    Returns:
        JsonSchema: Strict schema suitable for contexts such as sampling or
        external transports that require explicit constraints.

    """
    if schema == {}:
        return {"type": "object", "additionalProperties": False, "properties": {}, "required": []}

    clone = _clone_schema(schema)
    _enforce_strict_schema(clone, root=clone, path=())
    return clone


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _clone_schema(schema: JsonSchema) -> JsonSchema:
    """Return a deep copy of ``schema``.

    Returns:
        JsonSchema: Deep copy of the input schema.

    """
    if isinstance(schema, dict):
        return {k: _clone_schema(v) for k, v in schema.items()}
    if isinstance(schema, list):
        return [_clone_schema(item) for item in schema]
    return schema


def _describes_object(schema: Mapping[str, Any]) -> bool:
    """Return whether ``schema`` already encodes an object shape.

    Returns:
        bool: ``True`` when the schema includes object keywords.

    """
    if schema.get("type") == "object":
        return True
    return any(
        key in schema
        for key in ("properties", "patternProperties", "additionalProperties", "propertyNames", "dependentRequired")
    )


def _envelope_from_schema(schema: Mapping[str, Any], marker: str) -> SchemaEnvelope:
    """Construct an envelope from vendor metadata stored in ``schema``.

    Returns:
        SchemaEnvelope: Envelope reconstructed from raw schema.

    """
    wrap_field: str | None = None
    metadata = schema.get(marker)
    if isinstance(metadata, Mapping):
        wrap_field = str(metadata.get("field", DEFAULT_WRAP_FIELD))
    elif metadata:
        wrap_field = DEFAULT_WRAP_FIELD
    if wrap_field is not None:
        properties = schema.get("properties")
        if isinstance(properties, Mapping) and wrap_field not in properties:
            wrap_field = next(iter(properties.keys()), wrap_field)
    return SchemaEnvelope(schema=_clone_schema(schema), wrap_field=wrap_field)


def _strip_field(node: Any, field_name: str) -> None:
    """Remove ``field_name`` wherever it appears in ``node``.

    Args:
        node: Schema fragment to mutate.
        field_name: Property key to delete recursively.

    """
    if isinstance(node, MutableMapping):
        node.pop(field_name, None)
        for value in node.values():
            _strip_field(value, field_name)
    elif isinstance(node, list):
        for value in node:
            _strip_field(value, field_name)


def _relax_additional_properties(node: Any) -> None:
    """Drop ``additionalProperties: false`` from mapping nodes.

    Args:
        node: Schema fragment to mutate.

    """
    if isinstance(node, MutableMapping):
        if node.get("additionalProperties") is False:
            node.pop("additionalProperties")
        for value in node.values():
            _relax_additional_properties(value)
    elif isinstance(node, list):
        for value in node:
            _relax_additional_properties(value)


def _drop_top_level_property(schema: MutableMapping[str, Any], name: str) -> None:
    """Remove ``name`` from top-level ``properties`` and ``required`` lists.

    Args:
        schema: JSON Schema mapping to mutate in place.
        name: Property name to remove.

    """
    properties = schema.get("properties")
    if isinstance(properties, MutableMapping):
        properties.pop(name, None)
        if not properties:
            schema.pop("properties")

    required = schema.get("required")
    if isinstance(required, list) and name in required:
        required = [item for item in required if item != name]
        if required:
            schema["required"] = required
        else:
            schema.pop("required")


def _prune_empty_required(node: Any) -> None:
    """Delete empty ``required`` arrays produced by pruning.

    Args:
        node: Schema fragment to inspect recursively.

    """
    if isinstance(node, MutableMapping):
        required = node.get("required")
        if isinstance(required, list) and not required:
            node.pop("required")
        for value in node.values():
            _prune_empty_required(value)
    elif isinstance(node, list):
        for value in node:
            _prune_empty_required(value)


def _coerce_envelope(
    schema_like: Any,
    *,
    wrap_scalar: bool,
    wrap_field: str,
    mode: JsonSchemaMode,
    compress: bool,
    drop_titles: bool = True,
    relax_additional_properties: bool = True,
) -> SchemaEnvelope:
    """Convert ``schema_like`` into a :class:`SchemaEnvelope`.

    Args:
        schema_like: Mapping, :class:`SchemaEnvelope`, or Python annotation.
        wrap_scalar: Whether to box scalars.
        wrap_field: Synthetic property name used when boxing occurs.
        mode: JSON Schema mode supplied to :class:`pydantic.TypeAdapter`.
        compress: Whether to remove cosmetic metadata.
        drop_titles: Remove ``title`` fields during compression.
        relax_additional_properties: Drop ``additionalProperties: false``.

    Returns:
        SchemaEnvelope: Normalized schema envelope.

    Raises:
        SchemaError: If the schema cannot be derived.

    """
    if isinstance(schema_like, SchemaEnvelope):
        base_schema = (
            compress_schema(
                schema_like.schema, drop_titles=drop_titles, relax_additional_properties=relax_additional_properties
            )
            if compress
            else _clone_schema(schema_like.schema)
        )

        if schema_like.wrap_field is None and not _describes_object(base_schema):
            if not wrap_scalar:
                raise SchemaError(
                    "Schema describes a non-object value. Set wrap_scalar=True to comply with MCP output rules."
                )
            return ensure_object_schema(base_schema, wrap_scalar=True, wrap_field=wrap_field)

        return SchemaEnvelope(schema=base_schema, wrap_field=schema_like.wrap_field)

    if isinstance(schema_like, Mapping):
        base_schema = (
            compress_schema(
                schema_like, drop_titles=drop_titles, relax_additional_properties=relax_additional_properties
            )
            if compress
            else _clone_schema(schema_like)
        )
        return ensure_object_schema(base_schema, wrap_scalar=wrap_scalar, wrap_field=wrap_field)

    try:
        type_adapter = TypeAdapter(schema_like)
    except Exception as exc:
        raise SchemaError(f"Unable to create TypeAdapter for {schema_like!r}") from exc

    try:
        base_schema = type_adapter.json_schema(mode=mode)
    except Exception as exc:
        raise SchemaError(f"Unable to derive JSON schema for {schema_like!r}") from exc

    if compress:
        base_schema = compress_schema(
            base_schema, drop_titles=drop_titles, relax_additional_properties=relax_additional_properties
        )

    return ensure_object_schema(base_schema, wrap_scalar=wrap_scalar, wrap_field=wrap_field)


def _enforce_strict_schema(node: JsonSchema, *, root: JsonSchema, path: tuple[str, ...]) -> JsonSchema:
    """Recursively enforce strict JSON Schema semantics via depth-first search.

    Args:
        node: Schema fragment being inspected.
        root: Top-level schema for resolving references.
        path: Tuple representing the traversal path for diagnostics.

    Returns:
        JsonSchema: Strict schema node.

    Raises:
        SchemaError: If the schema fragment violates strict requirements.

    """
    if not isinstance(node, MutableMapping):
        raise SchemaError(f"Expected mapping for schema node; path={'/'.join(path)}")

    defs = node.get("$defs")
    if isinstance(defs, MutableMapping):
        for name, sub in defs.items():
            _enforce_strict_schema(sub, root=root, path=(*path, "$defs", name))

    definitions = node.get("definitions")
    if isinstance(definitions, MutableMapping):
        for name, sub in definitions.items():
            _enforce_strict_schema(sub, root=root, path=(*path, "definitions", name))

    type_hint = node.get("type")
    if type_hint == "object":
        if node.get("additionalProperties") is None:
            node["additionalProperties"] = False
        elif node.get("additionalProperties") not in (False,):
            raise SchemaError(f"Strict schema forbids additionalProperties other than False; path={'/'.join(path)}")

    properties = node.get("properties")
    if isinstance(properties, MutableMapping):
        node["required"] = list(properties.keys())
        for key, sub in properties.items():
            _enforce_strict_schema(sub, root=root, path=(*path, "properties", key))

    items = node.get("items")
    if isinstance(items, MutableMapping):
        node["items"] = _enforce_strict_schema(items, root=root, path=(*path, "items"))

    any_of = node.get("anyOf")
    if isinstance(any_of, list):
        node["anyOf"] = [
            _enforce_strict_schema(entry, root=root, path=(*path, "anyOf", str(i))) for i, entry in enumerate(any_of)
        ]

    one_of = node.get("oneOf")
    if isinstance(one_of, list):
        existing_any_of = node.get("anyOf")
        if not isinstance(existing_any_of, list):
            existing_any_of = []
        node["anyOf"] = existing_any_of + [
            _enforce_strict_schema(entry, root=root, path=(*path, "oneOf", str(i))) for i, entry in enumerate(one_of)
        ]
        node.pop("oneOf")

    all_of = node.get("allOf")
    if isinstance(all_of, list):
        if len(all_of) == 1:
            merged = _enforce_strict_schema(all_of[0], root=root, path=(*path, "allOf", "0"))
            node.pop("allOf")
            node.update({k: v for k, v in merged.items() if k != "$defs"})
        else:
            node["allOf"] = [
                _enforce_strict_schema(entry, root=root, path=(*path, "allOf", str(i)))
                for i, entry in enumerate(all_of)
            ]

    default = node.get("default", _MISSING)
    if default is None:
        node.pop("default", None)

    ref = node.get("$ref")
    if isinstance(ref, str) and len(node) > 1:
        resolved = _resolve_ref(root=root, ref=ref)
        if not isinstance(resolved, MutableMapping):
            raise SchemaError(f"Ref {ref!r} did not resolve to an object; path={'/'.join(path)}")
        node.update({k: v for k, v in resolved.items() if k != "$ref"})
        node.pop("$ref", None)
        return _enforce_strict_schema(node, root=root, path=path)

    return node


def _resolve_ref(*, root: Mapping[str, Any], ref: str) -> JsonSchema:
    """Resolve a ``$ref`` pointer against ``root``.

    Args:
        root: Schema root containing definitions.
        ref: JSON Pointer-style reference (``#/...``).

    Returns:
        JsonSchema: Referenced schema fragment.

    Raises:
        SchemaError: If the reference cannot be resolved to a mapping.

    """
    if not ref.startswith("#/"):
        raise SchemaError(f"Unexpected $ref format {ref!r}")

    target: JsonSchema | Mapping[str, Any] = root
    for part in ref[2:].split("/"):
        if not isinstance(target, Mapping) or part not in target:
            raise SchemaError(f"Unable to resolve $ref {ref!r} at {part!r}")
        target = target[part]  # type: ignore[index]
    if not isinstance(target, Mapping):
        raise SchemaError(f"Resolved $ref {ref!r} to non-mapping node {target!r}")
    return target
