import json
import time
from typing import List

import fire
import uvicorn
import tensorflow as tf
from fastapi import FastAPI
from fastapi import Request
from hypergol.utils import create_pydantic_type

from models.{{ modelName.asSnake }}.{{ modelName.asSnake }}_batch_processor import {{ modelName }}BatchProcessor
from data_models.{{ inputClass.asSnake }} import {{ inputClass }}
from data_models.{{ outputClass.asSnake }} import {{ outputClass }}

TITLE = 'Serve {{ modelName }}'
VERSION = '0.1'
DESCRIPTION = 'FastApi wrapper on {{ modelName }}, see /docs for API details'
USE_GPU = False
THREADS = None
MODEL_DIRECTORY = '<data directory>/<project>/<branch>/models/{{ modelName }}/<epoch>'


def load_model(modelDirectory, threads, useGPU):
    if not useGPU:
        tf.config.experimental.set_visible_devices([], 'GPU')
    if threads is not None:
        tf.config.threading.set_inter_op_parallelism_threads(threads)
        tf.config.threading.set_intra_op_parallelism_threads(threads)
    return tf.saved_model.load(export_dir=modelDirectory)


app = FastAPI(title=TITLE, version=VERSION, description=DESCRIPTION)
model = load_model(modelDirectory=MODEL_DIRECTORY, threads=THREADS, useGPU=USE_GPU)
batchProcessor = {{ modelName }}BatchProcessor(
    inputDataset=None,
    inputBatchSize=0,
    maxTokenCount=100,
    outputDataset=None
)
pyDantic{{ inputClass }} = create_pydantic_type({{ inputClass }})
pyDantic{{ outputClass }} = create_pydantic_type({{ outputClass }})


@app.middleware("http")
async def add_headers(request: Request, call_next):
    startTime = time.time()
    response = await call_next(request)
    response.headers["X-Model-Long-Name"] = model.get_long_name().numpy().decode('utf-8')
    response.headers["X-Process-Time"] = str(time.time() - startTime)
    return response


@app.get("/")
def test_main():
    return {
        'title': TITLE,
        'version': VERSION,
        'description': DESCRIPTION,
        'model': model.get_long_name().numpy().decode('utf-8')
    }


@app.post("/output", response_model=List[pyDantic{{ outputClass }}])
def get_outputs({{ inputClass.asPluralVariable }}: List[pyDantic{{ inputClass }}]):
    {{ inputClass.asPluralVariable }} = [{{ inputClass }}.from_data(json.loads({{ inputClass.asVariable }}.json())) for {{ inputClass.asVariable }} in {{ inputClass.asPluralVariable }}]
    tensorInputs = batchProcessor.process_input_batch({{ inputClass.asPluralVariable }})
    tensorOutputs = model.get_outputs(**tensorInputs)
    {{ outputClass.asPluralVariable }} = batchProcessor.process_output_batch(tensorOutputs)
    return [pyDantic{{ outputClass }}.parse_raw(json.dumps({{ outputClass.asVariable }}.to_data())) for {{ outputClass.asVariable }} in {{ outputClass.asPluralVariable }}]


def uvicorn_serve_{{ modelName.asSnake }}_run(port=8000, host='0.0.0.0'):
    uvicorn.run("serve_{{ modelName.asSnake }}:app", port=port, host=host, reload=True)


if __name__ == "__main__":
    tf.get_logger().setLevel('ERROR')
    fire.Fire(uvicorn_serve_{{ modelName.asSnake }}_run)
