DocForge2/src/docforge/command_reference.py

"""Deterministic command references derived from live CLI and MCP registrations."""

from __future__ import annotations

import argparse
import hashlib
import json
from collections.abc import Collection, Iterable, Mapping
from dataclasses import dataclass
from typing import Any, cast

from mcp.types import Tool


@dataclass(frozen=True)
class CliCommandReference:
    """One CLI command and the exact normalized usage emitted by argparse."""

    name: str
    invocation: str


@dataclass(frozen=True)
class McpToolReference:
    """One registered MCP tool with arguments derived from its input schema."""

    surface: str
    name: str
    required_arguments: tuple[str, ...]
    optional_arguments: tuple[str, ...]
    description: str
    input_schema_hash: str


def cli_command_references(
    parser: argparse.ArgumentParser | None = None,
) -> tuple[CliCommandReference, ...]:
    """Read CLI reference rows from the parser used by ``docforge``."""

    effective_parser = parser or _docforge_parser()
    commands = _command_parsers(effective_parser)
    return tuple(
        CliCommandReference(
            name=name,
            invocation=_normalize_usage(command_parser.format_usage()),
        )
        for name, command_parser in sorted(commands.items())
    )


def mcp_tool_references(
    tools: Iterable[Tool],
    *,
    expected_names: Collection[str] | None = None,
) -> tuple[McpToolReference, ...]:
    """Read MCP reference rows from registered tools returned by ``list_tools``.

    ``expected_names`` makes documentation generation fail closed when the selected
    server surface is incomplete or has drifted.
    """

    surface_by_name = _mcp_surface_by_name()
    references: list[McpToolReference] = []
    observed: set[str] = set()
    for tool in tools:
        if tool.name in observed:
            raise ValueError(f"MCP tool registration repeats {tool.name!r}")
        observed.add(tool.name)
        try:
            surface = surface_by_name[tool.name]
        except KeyError as error:
            raise ValueError(
                f"MCP tool {tool.name!r} has no declared capability surface"
            ) from error
        schema = tool.inputSchema
        properties = _schema_properties(schema)
        required = _schema_required(schema, properties)
        references.append(
            McpToolReference(
                surface=surface,
                name=tool.name,
                required_arguments=tuple(sorted(required)),
                optional_arguments=tuple(sorted(set(properties) - required)),
                description=_normalize_text(tool.description or ""),
                input_schema_hash=_canonical_hash(schema),
            )
        )
    if expected_names is not None:
        expected = set(expected_names)
        if observed != expected:
            missing = sorted(expected - observed)
            unexpected = sorted(observed - expected)
            raise ValueError(
                "Registered MCP tools do not match the requested reference surface: "
                f"missing={missing!r}, unexpected={unexpected!r}"
            )
    surface_rank = {"read": 0, "proposal": 1, "application": 2}
    return tuple(sorted(references, key=lambda item: (surface_rank[item.surface], item.name)))


def render_cli_reference_markdown(references: Iterable[CliCommandReference]) -> str:
    """Render a deterministic Markdown table for CLI commands."""

    rows = sorted(references, key=lambda item: item.name)
    lines = [
        "| Command | Invocation |",
        "|---|---|",
    ]
    lines.extend(
        f"| `{_escape_markdown(item.name)}` | `{_escape_markdown(item.invocation)}` |"
        for item in rows
    )
    return "\n".join(lines) + "\n"


def render_mcp_reference_markdown(references: Iterable[McpToolReference]) -> str:
    """Render a deterministic Markdown table for registered MCP tools."""

    surface_rank = {"read": 0, "proposal": 1, "application": 2}
    rows = sorted(references, key=lambda item: (surface_rank[item.surface], item.name))
    lines = [
        "| Surface | Tool | Required arguments | Optional arguments | "
        "Input schema SHA-256 | Description |",
        "|---|---|---|---|---|---|",
    ]
    lines.extend(
        "| "
        f"{_escape_markdown(item.surface)} | "
        f"`{_escape_markdown(item.name)}` | "
        f"{_argument_list(item.required_arguments)} | "
        f"{_argument_list(item.optional_arguments)} | "
        f"`{item.input_schema_hash}` | "
        f"{_escape_markdown(item.description) or '—'} |"
        for item in rows
    )
    return "\n".join(lines) + "\n"


def render_command_reference_markdown(
    cli_references: Iterable[CliCommandReference],
    mcp_references: Iterable[McpToolReference],
) -> str:
    """Render the complete deterministic CLI and MCP command reference."""

    return (
        "## CLI commands\n\n"
        f"{render_cli_reference_markdown(cli_references)}"
        "\n## MCP tools\n\n"
        f"{render_mcp_reference_markdown(mcp_references)}"
    )


def _docforge_parser() -> argparse.ArgumentParser:
    from .cli import _parser  # pyright: ignore[reportPrivateUsage]

    return _parser()


def _command_parsers(
    parser: argparse.ArgumentParser,
) -> Mapping[str, argparse.ArgumentParser]:
    actions = parser._actions  # pyright: ignore[reportPrivateUsage]
    subparsers = [
        action
        for action in actions
        if action.dest == "command" and isinstance(action.choices, dict)
    ]
    if len(subparsers) != 1:
        raise ValueError("CLI parser must define exactly one command subparser")
    return cast(Mapping[str, argparse.ArgumentParser], subparsers[0].choices)


def _mcp_surface_by_name() -> dict[str, str]:
    from .mcp_server import APPLICATION_TOOLS, PROPOSAL_TOOLS, READ_TOOLS

    groups = {
        "read": READ_TOOLS,
        "proposal": PROPOSAL_TOOLS,
        "application": APPLICATION_TOOLS,
    }
    surface_by_name: dict[str, str] = {}
    for surface, names in groups.items():
        for name in names:
            if name in surface_by_name:
                raise ValueError(f"MCP tool surface declaration repeats {name!r}")
            surface_by_name[name] = surface
    return surface_by_name


def _normalize_usage(usage: str) -> str:
    return _normalize_text(usage.removeprefix("usage: "))


def _normalize_text(value: str) -> str:
    return " ".join(value.split())


def _schema_properties(schema: Mapping[str, Any]) -> dict[str, Any]:
    raw_properties: object = schema.get("properties", {})
    if not isinstance(raw_properties, dict):
        raise ValueError("MCP tool input schema properties must be a string-keyed object")
    properties: dict[str, Any] = {}
    for name, value in cast(dict[object, object], raw_properties).items():
        if not isinstance(name, str):
            raise ValueError("MCP tool input schema properties must be a string-keyed object")
        properties[name] = value
    return properties


def _schema_required(schema: Mapping[str, Any], properties: Mapping[str, Any]) -> set[str]:
    raw_required: object = schema.get("required", [])
    if not isinstance(raw_required, list):
        raise ValueError("MCP tool input schema required arguments must be strings")
    required_names: set[str] = set()
    for name in cast(list[object], raw_required):
        if not isinstance(name, str):
            raise ValueError("MCP tool input schema required arguments must be strings")
        required_names.add(name)
    if not required_names <= set(properties):
        raise ValueError("MCP tool input schema requires an undeclared argument")
    return required_names


def _canonical_hash(payload: Mapping[str, Any]) -> str:
    encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
    return hashlib.sha256(encoded).hexdigest()


def _argument_list(arguments: tuple[str, ...]) -> str:
    if not arguments:
        return "—"
    return ", ".join(f"`{_escape_markdown(argument)}`" for argument in arguments)


def _escape_markdown(value: str) -> str:
    return value.replace("\\", "\\\\").replace("|", "\\|").replace("`", "\\`")