from typing import Callable

from flask import Flask, request, jsonify
from prometheus_client import generate_latest, CollectorRegistry, Summary

from image_prediction.utils import get_logger
from image_prediction.utils.process_wrapping import wrap_in_process

logger = get_logger()


def make_prediction_server(predict_fn: Callable):
    app = Flask(__name__)
    registry = CollectorRegistry(auto_describe=True)
    metric = Summary(
        f"redactmanager_imageClassification_seconds", f"Time spent on image-service classification.", registry=registry
    )

    @app.route("/ready", methods=["GET"])
    def ready():
        resp = jsonify("OK")
        resp.status_code = 200
        return resp

    @app.route("/health", methods=["GET"])
    def healthy():
        resp = jsonify("OK")
        resp.status_code = 200
        return resp

    def __failure():
        response = jsonify("Analysis failed")
        response.status_code = 500
        return response

    @app.route("/predict", methods=["POST"])
    @app.route("/", methods=["POST"])
    @metric.time()
    def predict():

        # Tensorflow does not free RAM. Workaround: Run prediction function (which instantiates a model) in sub-process.
        # See: https://stackoverflow.com/questions/39758094/clearing-tensorflow-gpu-memory-after-model-execution
        predict_fn_wrapped = wrap_in_process(predict_fn)

        logger.info("Analysing...")
        predictions = predict_fn_wrapped(request.data)

        if predictions is not None:
            response = jsonify(predictions)
            logger.info("Analysis completed.")
            return response
        else:
            logger.error("Analysis failed.")
            return __failure()

    @app.route("/prometheus", methods=["GET"])
    def prometheus():
        return generate_latest(registry=registry)

    return app