import os
import argparse
import traceback
import logging as log
import pandas as pd
import numpy as np
from flask import Flask, request, jsonify

from ergo.project import Project

prj         = None
classes     = None
num_outputs = 0
app         = Flask(__name__)

def get_input(req):
    try:
        # search in query parameters
        x = request.args.get('x')
        if x is not None and x != "":
            return x
    except:
        pass

    try:
        # search in form body
        x = request.form.get('x')
        if x is not None and x != "":
            return x
    except:
        pass


    try:
        # search as file upload
        x = request.files['x']
        if x is not None:
            return x
    except:
        pass

    try:
        # search for direct json body
        x = req.get_json(silent=True)
        if x is not None:
            return x
    except:
        pass

    return None

@app.route('/encode', methods=['POST', 'GET'])
def encode_route():
    global prj, classes, num_outputs

    try:
        xin = get_input(request)
        if xin is None:
            return "missing 'x' parameter", 400

        x = prj.logic.prepare_input(xin)

        return jsonify(list(x)), 200
    except Exception as e:
        return str(e), 400

@app.route('/', methods=['POST', 'GET'])
def infer_route():
    global prj, classes, num_outputs
    try:
        xin = get_input(request)
        if xin is None:
            return "missing 'x' parameter", 400

        # encode the input
        x = prj.logic.prepare_input(xin)
        # run inference
        y = prj.model.predict(np.array([x]))[0].tolist()
        # decode results
        num_y = len(y)
        resp = {}

        if num_y != num_outputs:
            return "expected %d output classes, got inference with %d results" % (num_outputs, num_y), 500

        resp = { classes[i] : y[i] for i in range(num_y) }

        return jsonify(resp), 200
    except Exception as e:
        return str(e), 400

def parse_args(argv):
    parser = argparse.ArgumentParser(prog="ergo serve", description="Load a model and expose an API that can be used to run its inference.",
            formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument("path", help="Path of the project.")

    parser.add_argument("-a", "--address", dest="address", action="store", type=str,
        help="IP address or hostname to bind to.")
    parser.add_argument("-p", "--port", dest="port", action="store", type=int, default=8080,
        help="TCP port to bind to.")
    parser.add_argument("--debug", dest="debug", action="store_true", default=False,
        help="Enable debug messages.")

    parser.add_argument("--profile", dest="profile", action="store_true", default=False,
        help="Enable profiler.")
    parser.add_argument("--prof-max", dest="restrictions", type=int, default=30,
        help="Maximum number of calls to profile.")

    parser.add_argument("--classes", dest="classes", default=None,
        help="Optional comma separated list of output classes.")

    args = parser.parse_args(argv)
    return args

def action_serve(argc, argv):
    global prj, app, classes, num_outputs

    args = parse_args(argv)
    prj = Project(args.path)
    err = prj.load()
    if err is not None:
        log.error("error while loading project: %s", err)
        quit()
    elif not prj.is_trained():
        log.error("no trained Keras model found for this project")
        quit()

    if args.classes is None:
        num_outputs = prj.model.output.shape[1]
        if prj.classes is None:
            classes = ["class_%d" % i for i in range(num_outputs)]
        else:
            classes = [prj.classes[i] for i in range(num_outputs)]
    else:
        classes = [s.strip() for s in args.classes.split(',') if s.strip() != ""]
        num_outputs = len(classes)

    if args.profile:
        from werkzeug.contrib.profiler import ProfilerMiddleware
        args.debug = True
        app.config['PROFILE'] = True
        app.wsgi_app = ProfilerMiddleware(app.wsgi_app, restrictions=[args.restrictions])

    app.run(host=args.address, port=args.port, debug=args.debug)
