@bentoml.service

svc.add_asgi_middleware()

from __future__ import annotations

import inspect
import json
import logging
import os
import traceback
import typing as t

import bentoml
from starlette.requests import Request
from starlette.responses import JSONResponse


def create_reroute_middleware(
    origin_route: str,
    destination_route: str,
    api_input: dict = {"input": ""},
):
    """Create reroute middleware

    Transforms a GET request from the origin route to a POST request to the destination
    route. This can be used to create or override GET routes, transforming them
    into POST requests against routes defined by the @bentoml.api decorator.

    For instance, we could intercept a call to /readyz and send a POST request
    to a custom route e.g. /myreadyz, which under the hood calls /predict.

    Parameters
    ----------
    origin_route
        Route that the user / client would call using a GET request
    destination_route
        Route to forward the request to using a POST request
    api_input
        Input to send to the destination API
    """
    class RerouteMiddleWare:
        def __init__(self, app):
            <http://self.app|self.app> = app

        async def __call__(self, scope, receive, send):
            async def _receive():
                await receive()
                message = {
                    "type": "http.request",
                    "body": json.dumps(api_input).encode(),
                }
                return message

            if "path" in scope and scope["method"] == "GET" and scope["path"] == origin_route:
                scope["path"] = destination_route
                scope["method"] = "POST"
                return await <http://self.app|self.app>(scope, _receive, send)

            return await <http://self.app|self.app>(scope, receive, send)

    return RerouteMiddleWare


def inject_myreadyz_route(api_input: dict, *, method_to_call: str = "predict"):
    """Inject a POST /myreadyz route

    This decorator should be used BEFORE the @bentoml.service decorator, to inject
    a route into the service.

    Parameters
    ----------
    api_input
        Input to send to `method_to_call`
    method_to_call
        Method to call, usually your primary predict method inside the service

    Example
    -------
    >>> @bentoml.service
        @inject_myreadyz_route(api_input={"input": [1, 2, 3]})
        class MyService:
            ...
    """
    def decorator(cls):
        m = getattr(cls, method_to_call)

        def handle_exception(e: Exception, ctx: bentoml.Context) -> str:
            logging.exception(e)
            ctx.response.status_code = 500
            return json.dumps(
                {
                    "status": "Not OK",
                    "error": repr(e),
                    "traceback": traceback.format_exc(),
                }
            )

        if inspect.iscoroutinefunction(m):
            async def myreadyz(self, input: str, ctx: bentoml.Context) -> str:
                try:
                    await getattr(self, method_to_call)(**api_input)
                    return json.dumps({"status": "OK"})
                except Exception as e:
                    return handle_exception(e=e, ctx=ctx)
        else:
            def myreadyz(self, input: str, ctx: bentoml.Context) -> str:
                try:
                    getattr(self, method_to_call)(**api_input)
                    return json.dumps({"status": "OK"})
                except Exception as e:
                    return handle_exception(e=e, ctx=ctx)

        setattr(cls, "myreadyz", bentoml.api(myreadyz, batchable=False))
        return cls

    return decorator


def inject_readyz_middleware(cls):
    """Inject readyz middleware

    Intercepts GET /readyz and sends them to POST /myreadyz.

    Should be used AFTER the @bentoml.service decorator

    Example
    -------
    >>> @inject_readyz_middleware
        @bentoml.service
        @inject_myreadyz_route(api_input={"input": [1, 2, 3]})
        class MyService:
            ...
    """
    ReadyzMiddleware = create_reroute_middleware(
        origin_route="/readyz",
        destination_route="/myreadyz",
        api_input={"input": ""},
    )
    cls.add_asgi_middleware(ReadyzMiddleware)
    return cls


@inject_readyz_middleware
@bentoml.service(resources={"cpu": "200m", "memory": "1Gi"}, traffic={"timeout": 10})
@inject_myreadyz_route(api_input={"input": "test"}, method_to_call="predict")
class MyService:
    @bentoml.api(batchable=False)
    def predict(self, input: str) -> list[str]:
        ...