"""Small, explicit primitives for exposing tools to an agent job.

The broker is a policy boundary, not a plugin system.  A process constructs a
``Registry`` itself, then creates one ``ToolBroker`` per job with an immutable
``JobPolicy``.  A broker snapshots the registered handlers, so later registry
changes cannot expand a running job's capabilities.
"""

from __future__ import annotations

import json
import math
import queue
import re
import threading
import time
from collections.abc import Callable, Mapping
from dataclasses import dataclass, field
from types import MappingProxyType
from urllib.parse import urlsplit

from .trace import OUTCOME_DENIED_OPERATOR, OUTCOME_DENIED_RUN, ToolTrace

_TOOL_ID_RE = re.compile(r"^[a-z][a-z0-9]*(?:[._-][a-z0-9]+)*$")
_VERSION_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
_HOST_LABEL_RE = re.compile(r"^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?$")


class ToolCallError(Exception):
    """A safe, expected tool failure that can be returned to the caller."""

    def __init__(self, code: str, message: str) -> None:
        super().__init__(message)
        self.code = code
        self.message = message


class ToolInputError(ToolCallError):
    """Raised by a tool when its arguments are invalid."""

    def __init__(self, message: str) -> None:
        super().__init__("invalid_input", message)


def _require_plain_int(name: str, value: object, minimum: int, maximum: int) -> int:
    if type(value) is not int:
        raise ValueError(f"{name} must be an integer (booleans are not accepted)")
    if not minimum <= value <= maximum:
        raise ValueError(f"{name} must be between {minimum} and {maximum}")
    return value


def _require_timeout(value: object) -> float:
    if isinstance(value, bool) or not isinstance(value, int | float):
        raise ValueError("timeout_seconds must be a finite number")
    timeout = float(value)
    if not math.isfinite(timeout) or not 0.001 <= timeout <= 300.0:
        raise ValueError("timeout_seconds must be between 0.001 and 300")
    return timeout


def _validate_tool_id(value: object) -> str:
    if not isinstance(value, str) or not 1 <= len(value) <= 96 or not _TOOL_ID_RE.fullmatch(value):
        raise ValueError("tool_id must be a lowercase, dotted identifier of at most 96 characters")
    return value


def _validate_version(value: object) -> str:
    if not isinstance(value, str) or not 1 <= len(value) <= 32 or not _VERSION_RE.fullmatch(value):
        raise ValueError("version must be a simple identifier of at most 32 characters")
    return value


def normalize_host(value: object) -> str:
    """Return a canonical ASCII hostname, rejecting URLs, ports and wildcards."""

    if not isinstance(value, str) or not value or len(value) > 253:
        raise ValueError("allowed hosts must be non-empty hostnames")
    if any(character.isspace() for character in value) or ":" in value or "/" in value or "*" in value:
        raise ValueError("allowed hosts must not contain schemes, ports, paths or wildcards")
    try:
        host = value.rstrip(".").encode("idna").decode("ascii").lower()
    except UnicodeError as exc:
        raise ValueError("allowed host is not a valid IDNA hostname") from exc
    if not host or any(not _HOST_LABEL_RE.fullmatch(label) for label in host.split(".")):
        raise ValueError("allowed host is malformed")
    return host


def _deep_freeze(value: object) -> object:
    if isinstance(value, Mapping):
        frozen = {str(key): _deep_freeze(item) for key, item in value.items()}
        return MappingProxyType(frozen)
    if isinstance(value, list | tuple):
        return tuple(_deep_freeze(item) for item in value)
    if isinstance(value, set):
        return frozenset(_deep_freeze(item) for item in value)
    return value


@dataclass(frozen=True)
class ToolRef:
    """An exact tool identity; versions are never resolved implicitly."""

    tool_id: str
    version: str

    def __post_init__(self) -> None:
        object.__setattr__(self, "tool_id", _validate_tool_id(self.tool_id))
        object.__setattr__(self, "version", _validate_version(self.version))


@dataclass(frozen=True)
class ToolSpec:
    """Immutable metadata published for a registered tool.

    ``origin`` says where the definition came from: ``"builtin"`` for a tool
    written in this repository, ``"mcp:<slug>"`` for one read from an MCP
    server.  It exists because ``description`` reaches the model's prompt, and
    a description written by a third party is untrusted input that a harness
    should fence the way it already fences tool *results*.  Provenance that is
    not carried alongside the text cannot be checked at the point of use.
    """

    tool_id: str
    version: str
    description: str
    input_schema: Mapping[str, object] = field(default_factory=dict)
    required_hosts: frozenset[str] = field(default_factory=frozenset)
    origin: str = "builtin"
    #: ``"tool"`` answers a question in-band; ``"skill"`` produces an artifact
    #: the run hands back to the user (a file, a plan, a document). Both travel
    #: the same registry, the same approval policy, the same broker and the
    #: same allowlist: the library treats them identically and the kind exists
    #: so surfaces can group them, never so policy can differ.
    kind: str = "tool"

    def __post_init__(self) -> None:
        object.__setattr__(self, "tool_id", _validate_tool_id(self.tool_id))
        object.__setattr__(self, "version", _validate_version(self.version))
        if self.kind not in {"tool", "skill"}:
            raise ValueError("kind must be 'tool' or 'skill'")
        if (
            not isinstance(self.description, str)
            or not self.description.strip()
            or len(self.description) > 2_000
        ):
            raise ValueError("description must be a non-empty string of at most 2,000 characters")
        if not isinstance(self.input_schema, Mapping):
            raise ValueError("input_schema must be a mapping")
        object.__setattr__(self, "input_schema", _deep_freeze(self.input_schema))
        object.__setattr__(
            self, "required_hosts", frozenset(normalize_host(host) for host in self.required_hosts)
        )
        if not isinstance(self.origin, str) or not self.origin.strip() or len(self.origin) > 128:
            raise ValueError("origin must be a short, non-empty string")
        object.__setattr__(self, "origin", self.origin.strip())

    @property
    def ref(self) -> ToolRef:
        return ToolRef(self.tool_id, self.version)

    @property
    def third_party(self) -> bool:
        """True when the description was written outside this repository."""

        return self.origin != "builtin"

    def to_dict(self) -> dict[str, object]:
        """Return JSON-ready metadata suitable for a model harness."""

        return {
            "id": self.tool_id,
            "version": self.version,
            "ref": f"{self.tool_id}@{self.version}",
            "description": self.description,
            "input_schema": _thaw_json(self.input_schema),
            "required_hosts": sorted(self.required_hosts),
            "origin": self.origin,
            "third_party": self.third_party,
            "kind": self.kind,
        }


@dataclass(frozen=True)
class ToolResult:
    """A JSON-safe result envelope returned for every broker call."""

    tool_id: str
    version: str
    ok: bool
    output: object | None = None
    error_code: str | None = None
    error_message: str | None = None
    elapsed_ms: int = 0

    @classmethod
    def success(cls, ref: ToolRef, output: object, elapsed_ms: int) -> ToolResult:
        return cls(ref.tool_id, ref.version, True, output=output, elapsed_ms=elapsed_ms)

    @classmethod
    def failure(cls, ref: ToolRef, code: str, message: str, elapsed_ms: int = 0) -> ToolResult:
        return cls(
            ref.tool_id,
            ref.version,
            False,
            error_code=code,
            error_message=message,
            elapsed_ms=elapsed_ms,
        )

    def to_dict(self) -> dict[str, object]:
        """Return a stable JSON-ready envelope for an inference harness."""

        value: dict[str, object] = {
            "tool_id": self.tool_id,
            "version": self.version,
            "ok": self.ok,
            "elapsed_ms": self.elapsed_ms,
        }
        if self.ok:
            value["output"] = _thaw_json(self.output)
        else:
            value["error"] = {
                "code": self.error_code,
                "message": self.error_message,
            }
        return value


@dataclass(frozen=True)
class JobPolicy:
    """Immutable limits for exactly one job's broker."""

    allowed_tools: frozenset[ToolRef]
    #: Kept in step with the server's per-run ceiling and the worker's hard
    #: cap; ``tests/test_tool_budget_agrees.py`` fails if they drift.
    max_calls: int = 10
    per_tool_quotas: Mapping[ToolRef, int] = field(default_factory=dict)
    timeout_seconds: float = 15.0
    max_input_bytes: int = 16_384
    max_output_bytes: int = 131_072
    allowed_hosts: frozenset[str] = field(default_factory=frozenset)

    def __post_init__(self) -> None:
        refs = frozenset(_coerce_ref(ref) for ref in self.allowed_tools)
        object.__setattr__(self, "allowed_tools", refs)
        object.__setattr__(self, "max_calls", _require_plain_int("max_calls", self.max_calls, 1, 100_000))
        object.__setattr__(self, "timeout_seconds", _require_timeout(self.timeout_seconds))
        object.__setattr__(
            self,
            "max_input_bytes",
            _require_plain_int("max_input_bytes", self.max_input_bytes, 1, 10 * 1024 * 1024),
        )
        object.__setattr__(
            self,
            "max_output_bytes",
            _require_plain_int("max_output_bytes", self.max_output_bytes, 1, 10 * 1024 * 1024),
        )
        quotas: dict[ToolRef, int] = {}
        for raw_ref, raw_limit in self.per_tool_quotas.items():
            ref = _coerce_ref(raw_ref)
            if ref not in refs:
                raise ValueError("per-tool quota refers to a tool that is not allowed")
            quotas[ref] = _require_plain_int("per-tool quota", raw_limit, 1, 100_000)
        object.__setattr__(self, "per_tool_quotas", MappingProxyType(quotas))
        object.__setattr__(
            self, "allowed_hosts", frozenset(normalize_host(host) for host in self.allowed_hosts)
        )


def _coerce_ref(value: object) -> ToolRef:
    if isinstance(value, ToolRef):
        return value
    if isinstance(value, tuple) and len(value) == 2:
        return ToolRef(value[0], value[1])
    raise ValueError("tool references must be ToolRef values or (tool_id, version) tuples")


@dataclass(frozen=True)
class ToolContext:
    """Capabilities and deadline supplied to a handler by the broker."""

    job_id: str
    deadline: float
    allowed_hosts: frozenset[str]

    def remaining_seconds(self) -> float:
        return max(0.0, self.deadline - time.monotonic())

    def require_url(self, url: str) -> str:
        """Validate an outbound HTTPS URL against this job's exact host allowlist.

        **This is a check a handler calls, not a boundary around a handler.**
        A handler runs on an ordinary thread in the agent process, with no
        import hook, no audit hook, no seccomp filter and no namespace, so a
        handler that imports ``socket`` and opens a connection is not stopped
        by anything here and does not appear in the trace as having done so.
        That was reproduced: a spec declaring no ``required_hosts``, therefore
        eligible for the ``local`` bundle, ran under a policy with an empty
        ``allowed_hosts`` and sent the prompt to a listener, and the trace
        recorded a clean success.

        What the framework does enforce is declaration: a tool that *declares*
        a host is refused unless the operator separately approved egress for
        it. What it cannot enforce is behaviour, which is why every tool in
        this repository is reviewed for what it imports, and why a third-party
        in-process tool is an operator's decision about code they are choosing
        to run in their own process rather than something this policy contains
        for them.

        Calling this is still worth doing. It is the difference between a tool
        that reaches one approved host and a tool that reaches any host, for
        every handler written in good faith, which is all of them here.
        """

        if not isinstance(url, str) or len(url) > 4_096:
            raise ToolCallError("host_not_allowed", "outbound URL is malformed")
        parsed = urlsplit(url)
        if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password:
            raise ToolCallError("host_not_allowed", "only HTTPS URLs without user information are allowed")
        try:
            host = normalize_host(parsed.hostname)
            port = parsed.port
        except (ValueError, UnicodeError) as exc:
            raise ToolCallError("host_not_allowed", "outbound URL is malformed") from exc
        if port not in (None, 443) or host not in self.allowed_hosts:
            raise ToolCallError("host_not_allowed", "outbound host is not allowed for this job")
        return host


ToolHandler = Callable[[Mapping[str, object], ToolContext], object]


@dataclass(frozen=True)
class _RegisteredTool:
    spec: ToolSpec
    handler: ToolHandler


class Registry:
    """A registry populated only by direct calls to ``register``.

    There is deliberately no entry-point discovery, import scanning, or dynamic
    module loading.
    """

    def __init__(self) -> None:
        self._tools: dict[ToolRef, _RegisteredTool] = {}
        self._lock = threading.RLock()

    def register(self, spec: ToolSpec, handler: ToolHandler) -> None:
        if not isinstance(spec, ToolSpec):
            raise TypeError("spec must be a ToolSpec")
        if not callable(handler):
            raise TypeError("handler must be callable")
        with self._lock:
            if spec.ref in self._tools:
                raise ValueError(f"tool {spec.tool_id} version {spec.version} is already registered")
            self._tools[spec.ref] = _RegisteredTool(spec, handler)

    def resolve(self, ref: ToolRef) -> _RegisteredTool:
        with self._lock:
            try:
                return self._tools[ref]
            except KeyError as exc:
                raise KeyError(f"tool {ref.tool_id} version {ref.version} is not registered") from exc

    def specs(self) -> tuple[ToolSpec, ...]:
        with self._lock:
            return tuple(
                entry.spec
                for _, entry in sorted(
                    self._tools.items(), key=lambda item: (item[0].tool_id, item[0].version)
                )
            )


def _copy_json(value: object, *, depth: int = 0, nodes: list | None = None) -> object:
    """Copy and validate the deliberately small JSON value subset used by tools."""

    if nodes is None:
        nodes = [0]
    nodes[0] += 1
    if nodes[0] > 10_000 or depth > 24:
        raise ValueError("JSON value is too complex")
    if value is None or type(value) is bool:
        return value
    if type(value) is str:
        return value
    if type(value) is int:
        if not -(2**63) <= value <= 2**63 - 1:
            raise ValueError("integers must fit in signed 64-bit range")
        return value
    if type(value) is float:
        if not math.isfinite(value):
            raise ValueError("non-finite numbers are not accepted")
        return value
    if isinstance(value, list | tuple):
        if len(value) > 10_000:
            raise ValueError("JSON arrays are too large")
        return [_copy_json(item, depth=depth + 1, nodes=nodes) for item in value]
    if isinstance(value, Mapping):
        if len(value) > 10_000:
            raise ValueError("JSON objects are too large")
        result: dict[str, object] = {}
        for key, item in value.items():
            if type(key) is not str or len(key) > 256 or any(ord(character) < 32 for character in key):
                raise ValueError("JSON object keys must be short strings without control characters")
            result[key] = _copy_json(item, depth=depth + 1, nodes=nodes)
        return result
    raise ValueError("tool values must contain only JSON-compatible types")


def _thaw_json(value: object) -> object:
    """Turn frozen metadata and validated output into ordinary JSON containers."""

    if isinstance(value, Mapping):
        return {key: _thaw_json(item) for key, item in value.items()}
    if isinstance(value, tuple | list | frozenset):
        return [_thaw_json(item) for item in value]
    return value


def _json_size(value: object) -> int:
    return len(json.dumps(value, ensure_ascii=False, allow_nan=False, separators=(",", ":")).encode("utf-8"))


class OperatorRefusal(PermissionError):
    """The worker's operator has not approved a tool this run asked for.

    Distinct from ``KeyError`` (not installed) and from ``tool_not_allowed``
    (not selected for this run) because the three have different remedies and
    collapsing them leaves the operator with nothing to act on.
    """

    def __init__(self, ref: ToolRef, reason: str) -> None:
        super().__init__(reason)
        self.ref = ref
        self.reason = reason


class ToolBroker:
    """Policy-enforcing, per-job dispatcher for registered tools.

    Two allowlists apply and they come from different parties.  ``policy``
    carries what the *server* selected for this run.  ``operator_policy``, if
    supplied, carries what the *worker's operator* has approved at all, and it
    is checked first: a server cannot reach a tool the operator withheld, even
    by naming it explicitly.
    """

    def __init__(
        self,
        registry: Registry,
        policy: JobPolicy,
        *,
        job_id: str,
        operator_policy: object = None,
    ) -> None:
        if not isinstance(registry, Registry):
            raise TypeError("registry must be a Registry")
        if not isinstance(policy, JobPolicy):
            raise TypeError("policy must be a JobPolicy")
        if not isinstance(job_id, str) or not job_id or len(job_id) > 128 or any(ord(c) < 32 for c in job_id):
            raise ValueError("job_id must be a non-empty, short string without control characters")

        snapshot: dict[ToolRef, _RegisteredTool] = {}
        for ref in policy.allowed_tools:
            snapshot[ref] = registry.resolve(ref)

        # Checked after resolution so the message can name the tool, and
        # before anything is stored so a refused run has no broker at all.
        if operator_policy is not None:
            refuse = getattr(operator_policy, "refusal_reason", None)
            if not callable(refuse):
                raise TypeError("operator_policy must expose refusal_reason(spec)")
            for ref in sorted(snapshot, key=lambda item: (item.tool_id, item.version)):
                reason = refuse(snapshot[ref].spec)
                if reason:
                    raise OperatorRefusal(ref, reason)

        self.job_id = job_id
        self.policy = policy
        self.operator_policy = operator_policy
        self._tools = MappingProxyType(snapshot)
        self._calls_started = 0
        self._calls_by_tool: dict[ToolRef, int] = {}
        self._quota_lock = threading.Lock()
        self.trace = ToolTrace(job_id=job_id)

    @property
    def calls_started(self) -> int:
        with self._quota_lock:
            return self._calls_started

    def tool_specs(self) -> tuple[ToolSpec, ...]:
        """Return only the exact tool versions granted to this job."""

        return tuple(
            self._tools[ref].spec
            for ref in sorted(self._tools, key=lambda item: (item.tool_id, item.version))
        )

    def tool_manifest(self) -> tuple[dict[str, object], ...]:
        """Return JSON-ready tool descriptions for prompt/tool-schema adapters."""

        return tuple(spec.to_dict() for spec in self.tool_specs())

    def tool_refs(self) -> tuple[str, ...]:
        """Return stable ``id@version`` capability identifiers."""

        return tuple(f"{spec.tool_id}@{spec.version}" for spec in self.tool_specs())

    def call(self, tool_id: str, version: str, arguments: Mapping[str, object]) -> ToolResult:
        """Call an exact tool version and return a non-throwing result envelope.

        Every outcome, including every refusal, is recorded on :attr:`trace`
        before the result is returned.  A denial the caller chooses to swallow
        is still a fact about the run.
        """

        result = self._dispatch(tool_id, version, arguments)
        self.trace.record_result(
            tool_id=result.tool_id,
            version=result.version,
            arguments=arguments,
            result_envelope=result.to_dict(),
        )
        return result

    def record_external_denial(
        self, tool_ref: str, reason: str, *, operator: bool = False
    ) -> None:
        """Record a refusal decided before the call reached this broker.

        An adapter that rejects an unknown tool id upstream would otherwise
        leave no trace of the attempt, and "the model tried to call a tool you
        did not select" is precisely what the user needs to be told.
        """

        tool_id, _, version = str(tool_ref).partition("@")
        self.trace.record(
            tool_id=tool_id,
            version=version,
            outcome=OUTCOME_DENIED_OPERATOR if operator else OUTCOME_DENIED_RUN,
            error_code="operator_not_approved" if operator else "tool_not_allowed",
            error_message=reason,
        )

    def _dispatch(self, tool_id: str, version: str, arguments: Mapping[str, object]) -> ToolResult:
        started = time.monotonic()
        try:
            ref = ToolRef(tool_id, version)
        except (TypeError, ValueError) as exc:
            fallback = ToolRef("invalid", "invalid")
            return ToolResult.failure(fallback, "invalid_tool_ref", str(exc))

        entry = self._tools.get(ref)
        if entry is None:
            return ToolResult.failure(ref, "tool_not_allowed", "tool id/version is not allowed for this job")

        missing_hosts = entry.spec.required_hosts.difference(self.policy.allowed_hosts)
        if missing_hosts:
            return ToolResult.failure(
                ref, "host_not_allowed", "the job policy does not allow this tool's network host"
            )

        try:
            copied_arguments = _copy_json(arguments)
            if not isinstance(copied_arguments, dict):
                raise ValueError("tool arguments must be a JSON object")
            if _json_size(copied_arguments) > self.policy.max_input_bytes:
                raise ValueError("tool arguments exceed the job input limit")
        except (TypeError, ValueError) as exc:
            return ToolResult.failure(ref, "invalid_input", str(exc))

        quota_error = self._reserve_call(ref)
        if quota_error is not None:
            return ToolResult.failure(ref, "quota_exceeded", quota_error)

        context = ToolContext(
            job_id=self.job_id,
            deadline=time.monotonic() + self.policy.timeout_seconds,
            allowed_hosts=self.policy.allowed_hosts,
        )
        completed: queue.Queue[tuple[bool, object]] = queue.Queue(maxsize=1)

        def run_handler() -> None:
            try:
                completed.put((True, entry.handler(copied_arguments, context)))
            except BaseException as exc:  # contained at the tool boundary
                completed.put((False, exc))

        worker = threading.Thread(
            target=run_handler,
            name=f"tool-{ref.tool_id}-{self.job_id}",
            daemon=True,
        )
        worker.start()
        try:
            succeeded, payload = completed.get(timeout=self.policy.timeout_seconds)
        except queue.Empty:
            return ToolResult.failure(
                ref,
                "timeout",
                f"tool exceeded its {self.policy.timeout_seconds:.3f} second deadline",
                _elapsed_ms(started),
            )

        if not succeeded:
            if isinstance(payload, ToolCallError):
                return ToolResult.failure(ref, payload.code, payload.message, _elapsed_ms(started))
            return ToolResult.failure(ref, "tool_error", "tool execution failed", _elapsed_ms(started))

        try:
            output = _copy_json(payload)
            if _json_size(output) > self.policy.max_output_bytes:
                return ToolResult.failure(
                    ref,
                    "output_too_large",
                    "tool output exceeds the job output limit",
                    _elapsed_ms(started),
                )
        except (TypeError, ValueError):
            return ToolResult.failure(
                ref,
                "invalid_output",
                "tool returned a value that is not valid bounded JSON",
                _elapsed_ms(started),
            )
        return ToolResult.success(ref, output, _elapsed_ms(started))

    def invoke(self, tool_id: str, version: str, arguments: Mapping[str, object]) -> ToolResult:
        """Model-harness entry point using an explicit id and version."""

        return self.call(tool_id, version, arguments)

    def invoke_ref(self, tool_ref: str, arguments: Mapping[str, object]) -> ToolResult:
        """Model-harness entry point using one exact ``id@version`` reference."""

        if not isinstance(tool_ref, str) or tool_ref.count("@") != 1:
            failure = ToolResult.failure(
                ToolRef("invalid", "invalid"),
                "invalid_tool_ref",
                "tool reference must use the exact id@version form",
            )
            self.trace.record_result(
                tool_id=failure.tool_id,
                version=failure.version,
                arguments=arguments,
                result_envelope=failure.to_dict(),
            )
            return failure
        tool_id, version = tool_ref.split("@", 1)
        return self.call(tool_id, version, arguments)

    def _reserve_call(self, ref: ToolRef) -> str | None:
        with self._quota_lock:
            if self._calls_started >= self.policy.max_calls:
                return "job call quota is exhausted"
            current = self._calls_by_tool.get(ref, 0)
            tool_limit = self.policy.per_tool_quotas.get(ref)
            if tool_limit is not None and current >= tool_limit:
                return "per-tool call quota is exhausted"
            self._calls_started += 1
            self._calls_by_tool[ref] = current + 1
            return None


def _elapsed_ms(started: float) -> int:
    return max(0, int(round((time.monotonic() - started) * 1_000)))