# app/main.py

def custom_openapi():
    """Customize OpenAPI Output with x-gram extensions for getgram MCP servers"""

    if app.openapi_schema:
        return app.openapi_schema

    openapi_schema = get_openapi(
        title=app.title,
        version=app.version,
        description=app.description,
        routes=app.routes,
        tags=app.openapi_tags,
    )

    # Add x-gram extensions to specific operations
    x_gram_extensions = {
        "search_documentation": {
            "x-gram": {
                "name": "search_django_docs",
                "summary": "Search Django documentation using semantic similarity",
                "description": """<context>
                This tool performs semantic search over Django 5.2.8 documentation using RAG (Retrieval-Augmented Generation).
                It returns relevant documentation chunks with similarity scores and token estimates for LLM context management.
                Perfect for finding specific Django functionality, code examples, and best practices.
                </context>

                <prerequisites>
                - Query should be natural language describing what you're looking for
                - Results are ranked by semantic similarity (score 0-1, higher is better)
                - Token estimates help manage LLM context windows
                - Supports filtering by minimum relevance score and maximum result count
                </prerequisites>""",
                "responseFilterType": "jq",
            }
        },
    }

    # Apply x-gram extensions to paths
    if "paths" in openapi_schema:
        for path, path_item in openapi_schema["paths"].items():
            for method, operation in path_item.items():
                if method.lower() in ["get", "post", "put", "delete", "patch"]:
                    operation_id = operation.get("operationId")
                    if operation_id in x_gram_extensions:
                        operation.update(x_gram_extensions[operation_id])

    app.openapi_schema = openapi_schema
    return app.openapi_schema