import logging
import tempfile
import tracemalloc

from flask import Flask, request, jsonify
from waitress import serve

from image_prediction.config import CONFIG
from image_prediction.predictor import Predictor, extract_image_metadata_pairs, classify_images
from image_prediction.response import build_response
from prometheus_client import Gauge, Counter
from prometheus_flask_exporter import PrometheusMetrics


def main():

    predictor = Predictor()
    logging.info("Predictor ready.")
    tracemalloc.start()

    app = Flask(__name__)
    metrics = PrometheusMetrics(app=app, path='/prometheus')

    file_counter = Counter("image_prediction_file_counter", "count processed files")
    ram_metric = Gauge("image_prediction_memory_usage", "Memory usage in Mb")

    @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

    @app.route("/", methods=["POST"])
    @metrics.summary('image_prediction_request_time_seconds', 'Time spent processing request')
    def predict():
        def do_monitoring():
            file_counter.inc()
            _, peak = tracemalloc.get_traced_memory()
            ram_metric.set(peak / 10 ** 6)

        pdf = request.data

        logging.debug("Running predictor on document...")
        with tempfile.NamedTemporaryFile() as tmp_file:
            tmp_file.write(pdf)
            image_metadata_pairs = extract_image_metadata_pairs(tmp_file.name)
            try:
                predictions, metadata = classify_images(predictor, image_metadata_pairs)
                if CONFIG.service.monitoring_enabled:
                    do_monitoring()
            except Exception as err:
                logging.warning("Analysis failed.")
                logging.exception(err)
                response = jsonify("Analysis failed.")
                response.status_code = 500
                return response
            logging.debug(f"Found images in document.")

        response = jsonify(build_response(list(predictions), list(metadata)))

        logging.info("Analysis completed.")
        return response

    run_prediction_server(app, mode=CONFIG.webserver.mode)


def run_prediction_server(app, mode="development"):
    if mode == "development":
        app.run(host=CONFIG.webserver.host, port=CONFIG.webserver.port, debug=True)
    elif mode == "production":
        serve(app, host=CONFIG.webserver.host, port=CONFIG.webserver.port)
    tracemalloc.stop()


if __name__ == "__main__":
    logging_level = CONFIG.service.logging_level
    logging.basicConfig(level=logging_level)
    logging.getLogger("flask").setLevel(logging.ERROR)
    logging.getLogger("urllib3").setLevel(logging.ERROR)
    logging.getLogger("werkzeug").setLevel(logging.ERROR)
    logging.getLogger("waitress").setLevel(logging.ERROR)

    main()