scripts/generate_aws_api_docs.py

# Regenerate server/providers/aws/*_api.md, server/actions/*.py, and openapi.
# yaml from *_service.py + botocore service models.

from __future__ import annotations

import ast
import json
import re
import sys
from pathlib import Path
from typing import Any

try:
    from botocore.model import Shape
    from botocore.session import Session
except ImportError:
    print("Install botocore: pip install botocore", file=sys.stderr)
    raise

ROOT = Path(__file__).resolve().parents[1]
AWS_DIR = ROOT / "server" / "providers" / "aws"
ACTIONS_DIR = ROOT / "server" / "actions"

GET_CLIENT_RE = re.compile(r'get_client\("([^"]+)"\)')
SKIP_METHODS = frozenset(
    {"__init__", "reset", "health", "describe_state", "handle"}
)

ACTION_CLASS_OVERRIDES: dict[str, str] = {
    "logs": "LogsAction",
    "apigateway": "ApiGatewayAction",
    "apigatewayv2": "ApiGatewayV2Action",
    "elasticache": "ElastiCacheAction",
    "cognito-idp": "CognitoIdpAction",
    "cognito-identity": "CognitoIdentityAction",
    "stepfunctions": "StepFunctionsAction",
    "secretsmanager": "SecretsManagerAction",
    "sesv2": "SesV2Action",
    "wafv2": "WafV2Action",
    "cloudwatch": "CloudWatchAction",
    "firehose": "FirehoseAction",
    "events": "EventBridgeAction",
    "route53": "Route53Action",
    "dynamodb": "DynamoDBAction",
    "sqs": "SQSAction",
    "sns": "SNSAction",
    "rds": "RDSAction",
    "ecs": "ECSAction",
    "ec2": "EC2Action",
    "glue": "GlueAction",
    "athena": "AthenaAction",
    "kinesis": "KinesisAction",
    "lambda": "LambdaAction",
    "iam": "IAMAction",
    "s3": "S3Action",
    "sts": "STSAction",
    "ses": "SESAction",
    "acm": "ACMAction",
    "ssm": "SSMAction",
}


def api_md_path(boto_id: str) -> Path:
    return AWS_DIR / f"{boto_id}_api.md"


def action_class_for(boto_id: str) -> str:
    if boto_id in ACTION_CLASS_OVERRIDES:
        return ACTION_CLASS_OVERRIDES[boto_id]
    parts = re.split(r"[-_]", boto_id.replace("-", "_"))
    pascal = "".join(
        p[:1].upper() + p[1:].lower() if p else "" for p in parts if p
    )
    return f"{pascal}Action"


def action_module_stem(boto_id: str) -> str:
    """Python import path (no .py); avoid reserved words."""
    stem = boto_id.replace("-", "_")
    if stem == "lambda":
        return "lambda_"
    return stem


def action_module_for(boto_id: str) -> str:
    return action_module_stem(boto_id) + ".py"


def parse_service_file(path: Path) -> tuple[str, list[str]]:
    text = path.read_text(encoding="utf-8")
    m = GET_CLIENT_RE.search(text)
    if not m:
        raise ValueError(f"No get_client() in {path}")
    boto_id = m.group(1)
    tree = ast.parse(text)
    ops: list[str] = []
    for node in tree.body:
        if isinstance(node, ast.ClassDef) and node.name.endswith("Service"):
            for item in node.body:
                if not isinstance(item, ast.FunctionDef):
                    continue
                name = item.name
                if name in SKIP_METHODS or not name[:1].isupper():
                    continue
                ops.append(name)
    return boto_id, ops


def snake_op(pascal: str) -> str:
    s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", pascal)
    return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower()


def example_for_shape(
    shape: Shape | None,
    *,
    depth: int = 0,
    max_depth: int = 8,
    budget: list[int],
) -> Any:
    """Build a JSON-serializable example; decrement budget for each scalar/struct key."""
    if shape is None or budget[0] <= 0:
        return None
    if depth > max_depth:
        return None

    t = shape.type_name

    if t == "string":
        budget[0] -= 1
        if getattr(shape, "enum", None):
            return shape.enum[0]
        name = shape.name or "value"
        return f"example-{name.lower()}"
    if t in ("integer", "long"):
        budget[0] -= 1
        return 0
    if t == "float" or t == "double":
        budget[0] -= 1
        return 0.0
    if t == "boolean":
        budget[0] -= 1
        return False
    if t == "timestamp":
        budget[0] -= 1
        return "2024-01-01T00:00:00Z"
    if t == "blob":
        budget[0] -= 1
        return "base64-encoded-bytes"
    if t == "map":
        budget[0] -= 1
        if shape.value and budget[0] > 0:
            v = example_for_shape(
                shape.value, depth=depth + 1, max_depth=max_depth, budget=budget
            )
            return {"exampleKey": v}
        return {}
    if t == "list":
        budget[0] -= 1
        if shape.member and budget[0] > 0:
            inner = example_for_shape(
                shape.member,
                depth=depth + 1,
                max_depth=max_depth,
                budget=budget,
            )
            return [inner] if inner is not None else []
        return []
    if t == "structure":
        out: dict[str, Any] = {}
        members = shape.members
        required = list(shape.required_members)
        for k in required:
            if budget[0] <= 0:
                break
            sub = members[k]
            if sub.type_name in ("structure", "list", "map") and depth >= 4:
                continue
            v = example_for_shape(
                sub, depth=depth + 1, max_depth=max_depth, budget=budget
            )
            if v is not None:
                out[k] = v
        return out

    budget[0] -= 1
    return None


def payload_and_meta(
    session: Session, boto_id: str, operation: str
) -> tuple[dict[str, Any], str, list[str]]:
    """Returns (example_payload, boto_method_snake, optional_param_names)."""
    try:
        model = session.get_service_model(boto_id)
    except Exception as e:
        return {}, snake_op(operation), [f"(botocore model unavailable: {e})"]

    try:
        op_model = model.operation_model(operation)
    except Exception:
        return {}, snake_op(operation), ["(unknown operation in botocore)"]

    inp = op_model.input_shape
    boto_method = snake_op(operation)
    optional_names: list[str] = []
    if inp is None:
        return {}, boto_method, optional_names

    budget = [35]
    payload = example_for_shape(inp, depth=0, max_depth=8, budget=budget)
    if not isinstance(payload, dict):
        payload = {}

    required = set(inp.required_members)
    for name in inp.members:
        if name not in required:
            optional_names.append(name)

    return payload, boto_method, optional_names


def curl_json_body(
    service: str, operation: str, payload: dict[str, Any]
) -> str:
    body = {
        "service": service,
        "operation": operation,
        "payload": payload,
    }
    return json.dumps(body, indent=2)


def write_api_md(
    boto_id: str,
    title: str,
    operations: list[str],
    session: Session,
) -> None:
    lines: list[str] = [
        f"# {title} Service API",
        "",
        "Requests use [`CloudEnvAction`](../../../models.py): `service`, `operation`, and `payload` (boto3 parameters for the operation).",
        "",
        f"Operations are enforced by IAM as `{boto_id}:<OperationName>` (see adapter).",
        "",
        "## Operations",
        "",
    ]
    for op in operations:
        payload, boto_method, optional_names = payload_and_meta(
            session, boto_id, op
        )
        lines.append(f"### {op}")
        lines.append("")
        lines.append(f"_Boto3 client method: `{boto_method}`_")
        lines.append("")
        if optional_names:
            preview = ", ".join(sorted(optional_names)[:25])
            more = (
                f" (+{len(optional_names) - 25} more)"
                if len(optional_names) > 25
                else ""
            )
            lines.append(
                f"**Other boto3 parameters (optional):** `{preview}`{more}"
            )
            lines.append("")
        json_body = curl_json_body(boto_id, op, payload)
        lines.append("```bash")
        lines.append("curl -X POST http://localhost:8000/step \\")
        lines.append('  -H "Content-Type: application/json" \\')
        lines.append("  -d @- <<'EOF'")
        lines.extend(json_body.split("\n"))
        lines.append("EOF")
        lines.append("```")
        lines.append("")
        lines.append("---")
        lines.append("")
    path = api_md_path(boto_id)
    path.write_text("\n".join(lines).rstrip() + "\n", encoding="utf-8")
    print(f"Wrote {path.relative_to(ROOT)}")


def write_action_py(boto_id: str, operations: list[str]) -> None:
    cls = action_class_for(boto_id)
    lines = [
        '"""Auto-generated from *_service.py — sync with botocore operation names."""',
        "from enum import Enum",
        "",
        "",
        f"class {cls}(str, Enum):",
    ]
    used: set[str] = set()
    for op in operations:
        key = re.sub(r"(.)([A-Z][a-z]+)", r"\1_\2", op)
        key = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", key).upper()
        key = re.sub(r"[^A-Z0-9]", "_", key)
        key = re.sub(r"_+", "_", key).strip("_")
        if not key or not key[0].isalpha():
            key = "OP_" + key
        base = key
        n = 2
        while key in used:
            key = f"{base}_{n}"
            n += 1
        used.add(key)
        lines.append(f'    {key} = "{op}"')
    path = ACTIONS_DIR / action_module_for(boto_id)
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    print(f"Wrote {path.relative_to(ROOT)}")


def write_openapi(
    session: Session, services: list[tuple[str, str, list[str]]]
) -> None:
    """services: (boto_id, title, operations)."""
    try:
        import yaml  # type: ignore
    except ImportError:
        print(
            "PyYAML required for openapi.yaml: pip install pyyaml",
            file=sys.stderr,
        )
        raise

    awsservice_values = sorted({b for b, _, _ in services})

    examples: dict[str, Any] = {}
    for boto_id, title, operations in sorted(services, key=lambda x: x[0]):
        for op in operations:
            payload, _, _ = payload_and_meta(session, boto_id, op)
            ex_key = f"{boto_id}__{op}"
            examples[ex_key] = {
                "summary": f"[{boto_id}] {op}",
                "description": f"Example {title} call via CloudEnv /step",
                "value": {
                    "service": boto_id,
                    "operation": op,
                    "payload": payload,
                },
            }

    doc: dict[str, Any] = {
        "openapi": "3.0.3",
        "info": {
            "title": "CloudEnv AWS Step API",
            "description": (
                "Single endpoint POST /step executes AWS operations via CloudEnvAction.\n"
                "Examples are named `{service}__{operation}` for filtering in Insomnia/Yaak.\n"
                "Each example sets service, operation, and payload (boto3 parameters).\n"
                "\n"
                "**Service groups:** use the `summary` prefix `[s3]`, `[lambda]`, etc., or search examples by service id."
            ),
            "version": "1.0.0",
        },
        "servers": [
            {
                "url": "http://localhost:8000",
                "description": "Local CloudEnv server",
            }
        ],
        "tags": [
            {"name": "CloudEnv", "description": "POST /step environment API"}
        ]
        + [
            {"name": f"aws.{bid}", "description": f"{title} (`{bid}`)"}
            for bid, title, _ in sorted(services, key=lambda x: x[0])
        ],
        "paths": {
            "/step": {
                "post": {
                    "operationId": "cloudenvStep",
                    "summary": "Execute a CloudEnv AWS action",
                    "description": (
                        "Body matches CloudEnvAction: service (AWS service id), "
                        "operation (PascalCase boto3 operation name), payload (parameters)."
                    ),
                    "tags": ["CloudEnv"],
                    "x-tagGroups": [
                        {
                            "name": "AWS services",
                            "tags": [
                                f"aws.{bid}" for bid, _, _ in sorted(services)
                            ],
                        }
                    ],
                    "requestBody": {
                        "required": True,
                        "content": {
                            "application/json": {
                                "schema": {
                                    "$ref": "#/components/schemas/CloudEnvAction"
                                },
                                "examples": examples,
                            }
                        },
                    },
                    "responses": {
                        "200": {
                            "description": "Observation from the environment",
                            "content": {
                                "application/json": {
                                    "schema": {
                                        "$ref": "#/components/schemas/CloudEnvObservation"
                                    }
                                }
                            },
                        },
                        "4XX": {"description": "Client / validation error"},
                        "5XX": {"description": "Server error"},
                    },
                }
            }
        },
        "components": {
            "schemas": {
                "CloudEnvAction": {
                    "type": "object",
                    "required": ["service", "operation"],
                    "properties": {
                        "service": {
                            "type": "string",
                            "description": "AWS service identifier (boto3 client name)",
                            "enum": awsservice_values,
                        },
                        "operation": {
                            "type": "string",
                            "description": "PascalCase operation (e.g. CreateBucket)",
                        },
                        "payload": {
                            "type": "object",
                            "additionalProperties": True,
                            "description": "Boto3 request parameters for the operation",
                        },
                        "intent": {
                            "type": "string",
                            "nullable": True,
                            "description": "Optional agent hint",
                        },
                    },
                },
                "CloudEnvObservation": {
                    "type": "object",
                    "properties": {
                        "services": {
                            "type": "object",
                            "additionalProperties": True,
                        },
                        "pipeline": {
                            "type": "object",
                            "additionalProperties": True,
                        },
                        "metrics": {
                            "type": "object",
                            "additionalProperties": True,
                        },
                        "operation_result": {
                            "type": "object",
                            "additionalProperties": True,
                            "description": "Serialized boto3 response from the last execute() call",
                        },
                        "reward": {"type": "number", "format": "float"},
                        "done": {"type": "boolean"},
                    },
                },
            }
        },
    }

    out = AWS_DIR / "openapi.yaml"
    out.write_text(
        yaml.dump(
            doc,
            sort_keys=False,
            allow_unicode=True,
            default_flow_style=False,
            width=120,
        ),
        encoding="utf-8",
    )
    print(f"Wrote {out.relative_to(ROOT)}")


def write_actions_init(service_rows: list[tuple[str, str, list[str]]]) -> None:
    """Generate server/actions/__init__.py re-exporting all *Action enums."""
    lines = [
        '"""Re-exports generated *Action enums (one per AWS service)."""',
    ]
    exports: list[str] = []
    for boto_id, _, ops in sorted(service_rows, key=lambda x: x[0]):
        mod = action_module_stem(boto_id)
        cls = action_class_for(boto_id)
        lines.append(f"from .{mod} import {cls}")
        exports.append(cls)
    lines.append("")
    lines.append("__all__ = [")
    for name in exports:
        lines.append(f'    "{name}",')
    lines.append("]")
    lines.append("")
    p = ACTIONS_DIR / "__init__.py"
    p.write_text("\n".join(lines), encoding="utf-8")
    print(f"Wrote {p.relative_to(ROOT)}")


def title_for(boto_id: str) -> str:
    special = {
        "logs": "CloudWatch Logs",
        "events": "EventBridge",
        "secretsmanager": "Secrets Manager",
        "sesv2": "SES v2",
        "apigateway": "API Gateway",
        "apigatewayv2": "API Gateway v2",
        "cognito-idp": "Cognito IDP",
        "cognito-identity": "Cognito Identity",
        "stepfunctions": "Step Functions",
        "wafv2": "WAFv2",
        "elasticache": "ElastiCache",
        "dynamodb": "DynamoDB",
        "cloudwatch": "CloudWatch",
    }
    if boto_id in special:
        return special[boto_id]
    return (
        boto_id.upper()
        if len(boto_id) <= 3
        else boto_id.replace("-", " ").title()
    )


def main() -> None:
    session = Session()
    service_files = sorted(AWS_DIR.glob("*_service.py"))
    rows: list[tuple[str, str, list[str]]] = []

    for sf in service_files:
        boto_id, ops = parse_service_file(sf)
        if not ops:
            print(f"Skip (no ops): {sf.name}", file=sys.stderr)
            continue
        title = title_for(boto_id)
        write_api_md(boto_id, title, ops, session)
        write_action_py(boto_id, ops)
        rows.append((boto_id, title, ops))

    write_actions_init(rows)
    write_openapi(session, rows)
    print("Done.")


if __name__ == "__main__":
    main()