"""Tool registry, and a JSON-Schema-subset validator for tool arguments.

The validator covers the subset a tool loop actually needs: ``type`` (including
the integer/number distinction, where ``1`` is a valid integer and ``True`` is
not), ``enum``, ``required``, ``additionalProperties``, ``items``, ``properties``
and the numeric/string bounds.  It is deliberately not a full JSON Schema
implementation -- ``allOf``/``$ref``/``format`` are absent, and claiming they
were supported would be a lie a caller eventually finds out.

``bool`` is checked before ``int`` because ``bool`` is a subclass of ``int`` in
Python: without that ordering ``True`` would validate as the integer ``1``.
"""

from __future__ import annotations

import re

__all__ = [
    "ToolError",
    "ToolNotFound",
    "ValidationError",
    "ToolCall",
    "ToolResult",
    "Tool",
    "ToolRegistry",
    "validate",
]


class ToolError(Exception):
    """Base class for every tool failure the loop can report."""


class ToolNotFound(ToolError):
    """The model named a tool that is not registered."""


class ValidationError(ToolError):
    """The model's arguments do not satisfy the tool's schema."""


def _is_type(instance, name):
    if name == "object":
        return isinstance(instance, dict)
    if name == "array":
        return isinstance(instance, list)
    if name == "string":
        return isinstance(instance, str)
    if name == "boolean":
        return isinstance(instance, bool)
    if name == "null":
        return instance is None
    if name == "integer":
        return isinstance(instance, int) and not isinstance(instance, bool)
    if name == "number":
        return isinstance(instance, (int, float)) and not isinstance(instance, bool)
    # An unknown type name is a schema bug, not a caller bug: fail loudly.
    raise ValidationError("unknown type %r in schema" % (name,))


def validate(instance, schema, path="$"):
    """Raise ``ValidationError`` unless ``instance`` satisfies ``schema``."""
    if schema is True or schema == {}:
        return
    if schema is False:
        raise ValidationError("%s: no value is allowed here" % path)
    if not isinstance(schema, dict):
        raise ValidationError(
            "%s: schema must be an object, got %s" % (path, type(schema).__name__)
        )

    if "enum" in schema:
        if instance not in schema["enum"]:
            raise ValidationError(
                "%s: %r is not one of %r" % (path, instance, schema["enum"])
            )
        return

    if "type" in schema:
        declared = schema["type"]
        names = declared if isinstance(declared, list) else [declared]
        if not any(_is_type(instance, n) for n in names):
            raise ValidationError(
                "%s: expected type %s, got %s" % (path, declared, type(instance).__name__)
            )

    if isinstance(instance, str):
        if "minLength" in schema and len(instance) < schema["minLength"]:
            raise ValidationError(
                "%s: shorter than minLength %d" % (path, schema["minLength"])
            )
        if "maxLength" in schema and len(instance) > schema["maxLength"]:
            raise ValidationError(
                "%s: longer than maxLength %d" % (path, schema["maxLength"])
            )
        if "pattern" in schema and re.search(schema["pattern"], instance) is None:
            raise ValidationError(
                "%s: %r does not match pattern %r" % (path, instance, schema["pattern"])
            )

    if isinstance(instance, bool):
        return

    if isinstance(instance, (int, float)):
        if "minimum" in schema and instance < schema["minimum"]:
            raise ValidationError("%s: %r is below minimum %r" % (path, instance, schema["minimum"]))
        if "maximum" in schema and instance > schema["maximum"]:
            raise ValidationError("%s: %r is above maximum %r" % (path, instance, schema["maximum"]))

    if isinstance(instance, list):
        if "minItems" in schema and len(instance) < schema["minItems"]:
            raise ValidationError(
                "%s: %d items is below minItems %d" % (path, len(instance), schema["minItems"])
            )
        if "maxItems" in schema and len(instance) > schema["maxItems"]:
            raise ValidationError(
                "%s: %d items is above maxItems %d" % (path, len(instance), schema["maxItems"])
            )
        if "items" in schema:
            for i, item in enumerate(instance):
                validate(item, schema["items"], "%s[%d]" % (path, i))

    if isinstance(instance, dict):
        for required in schema.get("required", []):
            if required not in instance:
                raise ValidationError("%s: missing required property %r" % (path, required))
        properties = schema.get("properties", {})
        for key, value in instance.items():
            if key in properties:
                validate(value, properties[key], "%s.%s" % (path, key))
            elif schema.get("additionalProperties") is False:
                raise ValidationError("%s: unexpected property %r" % (path, key))


class ToolCall(object):
    """One tool invocation requested by the model."""

    __slots__ = ("name", "arguments")

    def __init__(self, name, arguments=None):
        self.name = name
        self.arguments = dict(arguments) if arguments else {}

    def __repr__(self):
        return "ToolCall(%r, %r)" % (self.name, self.arguments)

    def __eq__(self, other):
        return (
            isinstance(other, ToolCall)
            and self.name == other.name
            and self.arguments == other.arguments
        )

    def as_dict(self):
        return {"name": self.name, "arguments": self.arguments}


class ToolResult(object):
    """The outcome of one tool invocation, as recorded in the transcript."""

    __slots__ = ("name", "ok", "value", "error")

    def __init__(self, name, ok, value=None, error=None):
        self.name = name
        self.ok = bool(ok)
        self.value = value
        self.error = error

    def __repr__(self):
        if self.ok:
            return "ToolResult(%r, ok=True, value=%r)" % (self.name, self.value)
        return "ToolResult(%r, ok=False, error=%r)" % (self.name, self.error)

    def __eq__(self, other):
        return (
            isinstance(other, ToolResult)
            and self.name == other.name
            and self.ok == other.ok
            and self.value == other.value
            and self.error == other.error
        )

    def as_dict(self):
        return {"name": self.name, "ok": self.ok, "value": self.value, "error": self.error}


class Tool(object):
    """A named, described, schema-checked callable."""

    def __init__(self, name, description, parameters, fn):
        if not isinstance(name, str) or not name:
            raise ValueError("tool name must be a non-empty string")
        if not callable(fn):
            raise ValueError("tool %r: fn must be callable" % name)
        if not isinstance(parameters, dict):
            raise ValueError("tool %r: parameters must be a JSON-schema object" % name)
        self.name = name
        self.description = description or ""
        self.parameters = parameters
        self.fn = fn

    def __call__(self, arguments=None):
        arguments = dict(arguments) if arguments else {}
        validate(arguments, self.parameters, "$.%s" % self.name)
        try:
            return self.fn(**arguments)
        except ToolError:
            raise
        except Exception as exc:
            raise ToolError(
                "tool %r raised %s: %s" % (self.name, type(exc).__name__, exc)
            )

    def manifest(self):
        return {
            "name": self.name,
            "description": self.description,
            "parameters": self.parameters,
        }

    def __repr__(self):
        return "Tool(%r)" % self.name


class ToolRegistry(object):
    """Name -> Tool, with duplicate rejection and a model-facing manifest."""

    def __init__(self, tools=()):
        self._tools = {}
        for tool in tools:
            self.add(tool)

    def add(self, tool):
        if not isinstance(tool, Tool):
            raise TypeError("expected Tool, got %s" % type(tool).__name__)
        if tool.name in self._tools:
            raise ValueError("duplicate tool name %r" % tool.name)
        self._tools[tool.name] = tool
        return tool

    def get(self, name):
        try:
            return self._tools[name]
        except KeyError:
            raise ToolNotFound("no tool named %r (have: %s)" % (name, ", ".join(sorted(self._tools))))

    def has(self, name):
        return name in self._tools

    def names(self):
        return sorted(self._tools)

    def manifest(self):
        return [self._tools[n].manifest() for n in sorted(self._tools)]

    def __len__(self):
        return len(self._tools)

    def __contains__(self, name):
        return name in self._tools