#!/usr/bin/env python3
"""Operator-owned tool MCP registry, validation, and Codex rendering.

Tool MCP servers expose third-party tools to Codex.  They are deliberately
separate from the MMO Agent MCP server (``mmo_mesh``), which launches and
supervises profile participants.
"""

from __future__ import annotations

import math
import os
import re
import shutil
from collections.abc import Mapping
from pathlib import Path
from typing import Any
from urllib.parse import urlsplit

from mmo_util import (
    config_root,
    parse_env_file,
    read_toml,
    valid_absolute_uri,
    valid_http_header_value,
    validate_id,
)
from mmo_version import MMO_SCHEMA_VERSION

TOOL_MCP_RESERVED_IDS = {"mmo_mesh"}
TOOL_MCP_TRANSPORTS = {"stdio", "streamable_http"}
TOOL_MCP_APPROVAL_MODES = {"auto", "prompt", "writes", "approve"}

_DOCUMENT_FIELDS = {"schema_version", "tool_mcp_servers"}
_COMMON_SERVER_FIELDS = {
    "transport",
    "enabled_tools",
    "default_tools_approval_mode",
    "startup_timeout_sec",
    "tool_timeout_sec",
    "supports_parallel_tool_calls",
    "tools",
}
_STDIO_FIELDS = {"command", "args", "env", "env_vars", "cwd"}
_HTTP_FIELDS = {
    "url",
    "bearer_token_env_var",
    "http_headers",
    "env_http_headers",
}
_OAUTH_FIELDS = {"auth", "scopes", "oauth", "oauth_resource", "environment_id"}
_REGISTRY_OWNED_POLICY_FIELDS = {"enabled", "required", "disabled_tools"}
_TOOL_FIELDS = {"approval_mode"}

_ENVIRONMENT_NAME = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
_HEADER_NAME = re.compile(r"[!#$%&'*+.^_`|~0-9A-Za-z-]+")
_SENSITIVE_ENVIRONMENT_NAME = re.compile(
    r"(?:^|_)(?:API_?KEY|KEY|TOKEN|SECRET|PASSWORD|PASSWD|CREDENTIAL|COOKIE|AUTH|PRIVATE)(?:_|$)",
    re.IGNORECASE,
)
_SENSITIVE_HEADER_NAME = re.compile(
    r"(?:^|-)(?:authorization|cookie|api-?key|key|token|secret|password|passwd|"
    r"credential|auth|private)(?:-|$)",
    re.IGNORECASE,
)


def tool_mcp_registry_root() -> Path:
    """Return the operator-owned directory containing registry fragments."""

    return config_root() / "tool-mcp.d"


def _reject_unknown_fields(value: Mapping[str, Any], allowed: set[str], label: str) -> None:
    unknown = sorted(set(value) - allowed)
    if unknown:
        raise ValueError(f"{label} has unknown fields: {', '.join(unknown)}")


def _nonempty_string(value: Any, label: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise ValueError(f"{label} must be a non-empty string")
    if any(ord(character) < 0x20 or ord(character) == 0x7F for character in value):
        raise ValueError(f"{label} cannot contain control characters")
    return value


def _string_list(value: Any, label: str, *, allow_empty: bool) -> list[str]:
    if not isinstance(value, list) or not all(isinstance(item, str) for item in value):
        raise ValueError(f"{label} must be an array of strings")
    if not allow_empty and not value:
        raise ValueError(f"{label} cannot be empty")
    if any(not item.strip() for item in value):
        raise ValueError(f"{label} cannot contain empty strings")
    if len(value) != len(set(value)):
        raise ValueError(f"{label} cannot contain duplicates")
    return list(value)


def _environment_name(value: Any, label: str) -> str:
    if not isinstance(value, str) or not _ENVIRONMENT_NAME.fullmatch(value):
        raise ValueError(f"{label} must be a valid environment variable name")
    return value


def _positive_seconds(value: Any, label: str) -> float:
    if (
        not isinstance(value, (int, float))
        or isinstance(value, bool)
        or not math.isfinite(float(value))
        or not 0 < float(value) <= 172_800
    ):
        raise ValueError(f"{label} must be a finite number between 0 and 172800")
    return float(value)


def _string_map(value: Any, label: str) -> dict[str, str]:
    if not isinstance(value, Mapping):
        raise ValueError(f"{label} must be a table of strings")
    result: dict[str, str] = {}
    for raw_key, raw_value in value.items():
        if not isinstance(raw_key, str) or not isinstance(raw_value, str):
            raise ValueError(f"{label} must be a table of strings")
        result[raw_key] = raw_value
    return result


def _validate_stdio_environment(value: Any, label: str) -> dict[str, str]:
    environment = _string_map(value, label)
    result: dict[str, str] = {}
    for name, raw_value in environment.items():
        _environment_name(name, f"{label} key")
        if _SENSITIVE_ENVIRONMENT_NAME.search(name):
            raise ValueError(
                f"{label}.{name} looks credential-bearing; use env_vars and credentials.env"
            )
        if "\x00" in raw_value:
            raise ValueError(f"{label}.{name} cannot contain NUL")
        result[name] = raw_value
    return dict(sorted(result.items()))


def _validate_headers(value: Any, label: str, *, environment_backed: bool) -> dict[str, str]:
    headers = _string_map(value, label)
    result: dict[str, str] = {}
    seen: set[str] = set()
    for raw_name, raw_value in headers.items():
        if not _HEADER_NAME.fullmatch(raw_name):
            raise ValueError(f"{label} contains an invalid HTTP header name: {raw_name!r}")
        normalized = raw_name.lower()
        if normalized in seen:
            raise ValueError(f"{label} contains duplicate case-insensitive header {raw_name!r}")
        seen.add(normalized)
        if environment_backed:
            _environment_name(raw_value, f"{label}.{raw_name}")
        else:
            if _SENSITIVE_HEADER_NAME.search(raw_name):
                raise ValueError(
                    f"{label}.{raw_name} is credential-bearing; use bearer_token_env_var "
                    "or env_http_headers"
                )
            if not valid_http_header_value(raw_value):
                raise ValueError(f"{label}.{raw_name} contains an invalid HTTP header value")
        result[raw_name] = raw_value
    return dict(sorted(result.items(), key=lambda item: item[0].lower()))


def _validate_http_url(value: Any, label: str) -> str:
    url = _nonempty_string(value, label)
    if not valid_absolute_uri(url):
        raise ValueError(f"{label} must be a valid HTTP(S) URL")
    try:
        parsed = urlsplit(url)
        _ = parsed.port
    except ValueError as exc:
        raise ValueError(f"{label} must be a valid HTTP(S) URL") from exc
    if parsed.scheme.lower() not in {"http", "https"} or not parsed.hostname:
        raise ValueError(f"{label} must be a valid HTTP(S) URL")
    if parsed.username is not None or parsed.password is not None:
        raise ValueError(f"{label} cannot contain embedded credentials")
    if parsed.fragment:
        raise ValueError(f"{label} cannot contain a URL fragment")
    return url


def _validate_stdio_command(value: Any, label: str) -> str:
    command = _nonempty_string(value, label)
    if "/" not in command and not command.startswith("~"):
        return command
    try:
        path = Path(command).expanduser()
    except RuntimeError as exc:
        raise ValueError(f"{label} has an unknown home-directory user") from exc
    if not path.is_absolute():
        raise ValueError(f"{label} must be a PATH executable name or an absolute path")
    return str(path.resolve())


def _validate_tool_policy(
    value: Any,
    *,
    label: str,
    enabled_tools: set[str],
) -> dict[str, dict[str, str]]:
    if value is None:
        return {}
    if not isinstance(value, Mapping):
        raise ValueError(f"{label} must be a table")
    result: dict[str, dict[str, str]] = {}
    for raw_name, raw_policy in value.items():
        tool_name = _nonempty_string(raw_name, f"{label} tool name")
        if tool_name not in enabled_tools:
            raise ValueError(f"{label}.{tool_name} is not present in enabled_tools")
        if not isinstance(raw_policy, Mapping):
            raise ValueError(f"{label}.{tool_name} must be a table")
        _reject_unknown_fields(raw_policy, _TOOL_FIELDS, f"{label}.{tool_name}")
        approval = raw_policy.get("approval_mode")
        if not isinstance(approval, str) or approval not in TOOL_MCP_APPROVAL_MODES:
            raise ValueError(
                f"{label}.{tool_name}.approval_mode must be one of "
                f"{sorted(TOOL_MCP_APPROVAL_MODES)}"
            )
        result[tool_name] = {"approval_mode": approval}
    return dict(sorted(result.items()))


def validate_tool_mcp_server(server_id: str, value: Any, *, label: str) -> dict[str, Any]:
    """Validate and normalize one operator-defined tool MCP server."""

    server_id = validate_id(server_id, "tool MCP server id")
    if server_id in TOOL_MCP_RESERVED_IDS:
        raise ValueError(f"tool MCP server id {server_id!r} is reserved for Agent MCP")
    if not isinstance(value, Mapping):
        raise ValueError(f"{label} must be a table")
    if "bearer_token" in value:
        raise ValueError(f"{label}.bearer_token is forbidden; use bearer_token_env_var")
    oauth_fields = sorted(set(value) & _OAUTH_FIELDS)
    if oauth_fields:
        raise ValueError(
            f"{label} uses unsupported OAuth fields: {', '.join(oauth_fields)}; "
            "the Tool MCP registry supports environment-backed authentication only"
        )
    policy_fields = sorted(set(value) & _REGISTRY_OWNED_POLICY_FIELDS)
    if policy_fields:
        raise ValueError(f"{label} cannot set agent-owned fields: {', '.join(policy_fields)}")
    transport = value.get("transport")
    if not isinstance(transport, str) or transport not in TOOL_MCP_TRANSPORTS:
        raise ValueError(f"{label}.transport must be one of {sorted(TOOL_MCP_TRANSPORTS)}")
    allowed = _COMMON_SERVER_FIELDS | (_STDIO_FIELDS if transport == "stdio" else _HTTP_FIELDS)
    _reject_unknown_fields(value, allowed, label)

    enabled_tools = _string_list(
        value.get("enabled_tools"), f"{label}.enabled_tools", allow_empty=False
    )
    approval = value.get("default_tools_approval_mode")
    if not isinstance(approval, str) or approval not in TOOL_MCP_APPROVAL_MODES:
        raise ValueError(
            f"{label}.default_tools_approval_mode must be one of {sorted(TOOL_MCP_APPROVAL_MODES)}"
        )
    normalized: dict[str, Any] = {
        "transport": transport,
        "enabled_tools": enabled_tools,
        "default_tools_approval_mode": approval,
        "supports_parallel_tool_calls": False,
    }
    if "supports_parallel_tool_calls" in value:
        if not isinstance(value["supports_parallel_tool_calls"], bool):
            raise ValueError(f"{label}.supports_parallel_tool_calls must be boolean")
        normalized["supports_parallel_tool_calls"] = value["supports_parallel_tool_calls"]
    for timeout_field in ("startup_timeout_sec", "tool_timeout_sec"):
        if timeout_field in value:
            normalized[timeout_field] = _positive_seconds(
                value[timeout_field], f"{label}.{timeout_field}"
            )
    normalized["tools"] = _validate_tool_policy(
        value.get("tools"),
        label=f"{label}.tools",
        enabled_tools=set(enabled_tools),
    )

    if transport == "stdio":
        normalized["command"] = _validate_stdio_command(value.get("command"), f"{label}.command")
        normalized["args"] = _string_list(value.get("args", []), f"{label}.args", allow_empty=True)
        if any("\x00" in argument for argument in normalized["args"]):
            raise ValueError(f"{label}.args cannot contain NUL")
        normalized["env"] = _validate_stdio_environment(value.get("env", {}), f"{label}.env")
        env_vars = _string_list(value.get("env_vars", []), f"{label}.env_vars", allow_empty=True)
        for name in env_vars:
            _environment_name(name, f"{label}.env_vars")
        overlap = sorted(set(normalized["env"]) & set(env_vars))
        if overlap:
            raise ValueError(
                f"{label} defines the same variables in env and env_vars: {', '.join(overlap)}"
            )
        normalized["env_vars"] = env_vars
        if "cwd" in value:
            raw_cwd = _nonempty_string(value["cwd"], f"{label}.cwd")
            try:
                cwd = Path(raw_cwd).expanduser()
            except RuntimeError as exc:
                raise ValueError(f"{label}.cwd has an unknown home-directory user") from exc
            if not cwd.is_absolute():
                raise ValueError(f"{label}.cwd must expand to an absolute path")
            normalized["cwd"] = str(cwd.resolve())
    else:
        normalized["url"] = _validate_http_url(value.get("url"), f"{label}.url")
        if "bearer_token_env_var" in value:
            normalized["bearer_token_env_var"] = _environment_name(
                value["bearer_token_env_var"], f"{label}.bearer_token_env_var"
            )
        normalized["http_headers"] = _validate_headers(
            value.get("http_headers", {}),
            f"{label}.http_headers",
            environment_backed=False,
        )
        normalized["env_http_headers"] = _validate_headers(
            value.get("env_http_headers", {}),
            f"{label}.env_http_headers",
            environment_backed=True,
        )
        static_names = {name.lower() for name in normalized["http_headers"]}
        environment_names = {name.lower() for name in normalized["env_http_headers"]}
        overlap = sorted(static_names & environment_names)
        if overlap:
            raise ValueError(
                f"{label} defines headers in both http_headers and env_http_headers: "
                + ", ".join(overlap)
            )
    return normalized


def _load_tool_mcp_fragment(path: Path) -> dict[str, dict[str, Any]]:
    data = read_toml(path)
    _reject_unknown_fields(data, _DOCUMENT_FIELDS, f"tool MCP registry {path}")
    schema = data.get("schema_version")
    if not isinstance(schema, int) or isinstance(schema, bool) or schema != MMO_SCHEMA_VERSION:
        raise ValueError(f"unsupported tool MCP registry schema_version in {path}")
    raw_servers = data.get("tool_mcp_servers", {})
    if not isinstance(raw_servers, Mapping):
        raise ValueError(f"tool_mcp_servers must be a table in {path}")
    return {
        server_id: validate_tool_mcp_server(
            server_id,
            server,
            label=f"tool MCP server {server_id} in {path}",
        )
        for server_id, server in raw_servers.items()
    }


def load_tool_mcp_registry_with_sources() -> tuple[dict[str, dict[str, Any]], dict[str, str]]:
    """Load deterministic operator fragments, replacing duplicate entries whole."""

    registry: dict[str, dict[str, Any]] = {}
    sources: dict[str, str] = {}
    root = tool_mcp_registry_root()
    if root.is_dir():
        for path in sorted(root.glob("*.toml")):
            for server_id, server in _load_tool_mcp_fragment(path).items():
                registry[server_id] = server
                sources[server_id] = str(path)
    return dict(sorted(registry.items())), dict(sorted(sources.items()))


def load_tool_mcp_registry() -> dict[str, dict[str, Any]]:
    return load_tool_mcp_registry_with_sources()[0]


def validate_tool_mcp_grants(
    value: Any,
    registry: Mapping[str, Mapping[str, Any]],
    *,
    label: str,
) -> dict[str, dict[str, Any]]:
    """Validate an agent's required/optional grants against operator maxima."""

    if value is None:
        return {}
    if not isinstance(value, Mapping):
        raise ValueError(f"{label} must be a table")
    result: dict[str, dict[str, Any]] = {}
    for raw_server_id, raw_grant in value.items():
        server_id = validate_id(raw_server_id, "tool MCP server id")
        if server_id not in registry:
            raise ValueError(
                f"{label}.{server_id} references an undefined operator tool MCP server; "
                f"define it under {tool_mcp_registry_root()}"
            )
        if not isinstance(raw_grant, Mapping):
            raise ValueError(f"{label}.{server_id} must be a table")
        _reject_unknown_fields(
            raw_grant,
            {"required", "enabled_tools"},
            f"{label}.{server_id}",
        )
        required = raw_grant.get("required", True)
        if not isinstance(required, bool):
            raise ValueError(f"{label}.{server_id}.required must be boolean")
        maximum = list(registry[server_id]["enabled_tools"])
        selected = (
            maximum
            if "enabled_tools" not in raw_grant
            else _string_list(
                raw_grant["enabled_tools"],
                f"{label}.{server_id}.enabled_tools",
                allow_empty=False,
            )
        )
        outside_maximum = sorted(set(selected) - set(maximum))
        if outside_maximum:
            raise ValueError(
                f"{label}.{server_id}.enabled_tools exceeds the operator allowlist: "
                + ", ".join(outside_maximum)
            )
        result[server_id] = {"required": required, "enabled_tools": selected}
    return dict(sorted(result.items()))


def tool_mcp_environment_names(server: Mapping[str, Any]) -> list[str]:
    """Return process environment names referenced by a normalized definition."""

    names: set[str] = set()
    if server["transport"] == "stdio":
        names.update(str(name) for name in server.get("env_vars", []))
    else:
        names.update(tool_mcp_http_environment_names(server))
    return sorted(names)


def tool_mcp_http_environment_names(server: Mapping[str, Any]) -> list[str]:
    """Return environment names whose values become HTTP header material."""

    if server["transport"] != "streamable_http":
        return []
    names = {str(name) for name in server.get("env_http_headers", {}).values()}
    bearer = server.get("bearer_token_env_var")
    if bearer:
        names.add(str(bearer))
    return sorted(names)


def codex_tool_mcp_server_config(
    server: Mapping[str, Any],
    grant: Mapping[str, Any] | None,
) -> dict[str, Any]:
    """Compile one normalized definition/grant to Codex's mcp_servers shape."""

    transport = str(server["transport"])
    fields = _STDIO_FIELDS if transport == "stdio" else _HTTP_FIELDS
    result = {
        key: server[key] for key in sorted(fields) if key in server and server[key] not in ({}, [])
    }
    maximum = list(server["enabled_tools"])
    selected = set(grant["enabled_tools"] if grant is not None else ())
    result.update(
        {
            "enabled": grant is not None,
            "required": bool(grant["required"]) if grant is not None else False,
            "supports_parallel_tool_calls": bool(server["supports_parallel_tool_calls"]),
            "enabled_tools": maximum,
            # Always emit this array so a native role can clear a narrower
            # deny-list inherited from its parent config layer.
            "disabled_tools": [tool for tool in maximum if tool not in selected],
            "default_tools_approval_mode": server["default_tools_approval_mode"],
        }
    )
    for timeout_field in ("startup_timeout_sec", "tool_timeout_sec"):
        if timeout_field in server:
            result[timeout_field] = server[timeout_field]
    if server.get("tools"):
        result["tools"] = server["tools"]
    return result


def tool_mcp_readiness(
    server_id: str,
    server: Mapping[str, Any],
    *,
    sources: Mapping[str, str] | None = None,
) -> dict[str, Any]:
    """Return non-launching command, cwd, and environment readiness."""

    values = parse_env_file(config_root() / "credentials.env")
    values.update({key: value for key, value in os.environ.items() if value})
    http_environment = set(tool_mcp_http_environment_names(server))
    environment: dict[str, bool] = {}
    invalid_environment: list[str] = []
    for name in tool_mcp_environment_names(server):
        value = values.get(name)
        valid = bool(value) and "\x00" not in str(value)
        if valid and name in http_environment:
            valid = valid_http_header_value(str(value))
        environment[name] = valid
        if value and not valid:
            invalid_environment.append(name)
    result: dict[str, Any] = {
        "id": server_id,
        "source": (sources or {}).get(server_id),
        "transport": server["transport"],
        "enabled_tools": list(server["enabled_tools"]),
        "default_tools_approval_mode": server["default_tools_approval_mode"],
        "environment": environment,
        "invalid_environment": invalid_environment,
    }
    if server["transport"] == "stdio":
        command = str(server["command"])
        command_path = Path(command).expanduser()
        resolved = (
            str(command_path.resolve())
            if "/" in command and command_path.is_file() and os.access(command_path, os.X_OK)
            else shutil.which(command)
        )
        result["command"] = command
        result["resolved_command"] = resolved
        cwd = server.get("cwd")
        result["cwd"] = cwd
        result["cwd_ready"] = cwd is None or Path(str(cwd)).is_dir()
        result["transport_ready"] = bool(resolved) and bool(result["cwd_ready"])
    else:
        result["url"] = server["url"]
        result["transport_ready"] = True
    result["environment_ready"] = all(environment.values())
    result["ready"] = bool(result["transport_ready"]) and bool(result["environment_ready"])
    return result