# Copyright 2019 The Vearch Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# ==============================================================================

"""``main.py`` is the entry of programs, you can run ``python main.py``.
Options:
    --port        Define the port your server run on
    --gpu         Define the GPU your server run on
    --model_name  Define the model you need
"""

import os
import json
import queue
import time
import signal
import requests
import threading
import shortuuid
import argparse
import traceback
import concurrent.futures
from multiprocessing import Queue, Process
from asyncio.futures import wrap_future
import tornado
import tornado.web
import tornado.httpserver
from tornado.httputil import HTTPServerRequest
from typing import Dict, List, Union
from tornado.options import define, options

import config
import util
import controller

_INSERT = '_insert'
_SEARCH = '_search'
_HEADERS = {'content-type': 'application/json'}
# InvalidURL = requests.InvalidURL


class PackageProcess(Process):
    """The Process dealing request and return result."""

    def __init__(self, 
                 gpu_id: str,
                 model_name: str,
                 input_queue: Queue,
                 output_queue: Queue,
                 debug: bool = False
                 ) -> None:
        Process.__init__(self)
        self.gpu_id = gpu_id               # The GPU id of model run on
        self.model_name = model_name       # image_retrieval or face_retrieval
        self.model = None
        self.debug = debug
        self.input_queue = input_queue
        self.output_queue = output_queue
        self.watch_queue = queue.Queue()   # thread queue for performance

    def build(self):
        """Initialization process environment.
            1.set visible GPU;
            2.load defined model
        """
        os.environ['CUDA_VISIBLE_DEVICES'] = self.gpu_id
        self.daemon_thread()
        self.model = controller.load_model(self.model_name)

    @staticmethod
    def daemon_thread():
        """Define a daemon thread listen the status of parent process."""
        def kill():
            while True:
                if os.getppid() == 1:
                    os.kill(os.getpid(), signal.SIGKILL)
                time.sleep(1)
        t = threading.Thread(target=kill)
        t.start()

    @staticmethod
    def judge(uri_list: List[str], request: 'Request') -> None:
        """Raise a Exception when create a db that name in [_cluster, list, db, space]
        :param uri_list: The list of uri.
        :param request: The dict of request body.
        :return None
        :raise Exception
        """
        if uri_list[1] == 'db' and uri_list[2] == '_create':
            if not request.body:
                raise requests.RequestException('The request data can not be empty!')
            data = json.loads(request.body, encoding='utf8')
            if 'name' not in data:
                raise requests.RequestException('The name of db is required in request data!')
            if data['name'] in ['_cluster', 'list', 'db', 'space']:
                raise requests.RequestException('The name of db can not in [_cluster, list, db, space]!')

    def deal(self, uuid: str, request: 'Request') -> None:
        """Deal the request.
        Repackage the request body and send it to VectorDB and put the response into `response_queue`
        :param uuid: The uuid of request.
        :param request: The request object.
        :return None
        """
        try:
            uri_list = request.uri.split('?')[0].split('/')
            if len(uri_list) < 3:
                raise requests.RequestException('Bad Request, Page not Found.')
            elif uri_list[1] in ['_cluster', 'list', 'db', 'space']:
                self.judge(uri_list, request)
                ip = f'{config.master_address}{request.uri}'
            elif request.method == 'POST':
                self.post(uri_list, request)
                ip = f'{config.router_address}{request.uri}'
            else:
                ip = f'{config.router_address}{request.uri}'
            res = requests.request(request.method, ip, data=request.body, headers=request.headers)
            result = Response(res)
            self.output_queue.put((uuid, True, result))
        except Exception as err:
            if self.debug:
                traceback.print_exc()
            self.output_queue.put((uuid, False, str(err)))

    def post(self, uri_list: List[str], request: 'Request') -> None:
        """Extract feature in request if necessary.
        :param uri_list: The list of uri.
        :param request: The dict of request body.
        :return None
        :raise NotImplementedError
        """
        flag = False
        operate = uri_list[3] if len(uri_list) >= 4 else None
        data = json.loads(request.body, encoding='utf8')
        if operate == _SEARCH:
            for d in data['query']['sum']:
                if not isinstance(d['feature'], list):
                    flag = True
                    d['feature'] = util.normlize(self.model.encode(d['feature']))
        elif operate == '_msearch':
            raise NotImplementedError('This API do not support msearch operate')
        else:
            for key in data:
                if isinstance(data[key], dict) and 'feature' in data[key]:
                    field = data[key]['feature']
                    if not isinstance(field, list):
                        flag = True
                        data[key]['feature'] = util.normlize(self.model.encode(field))
        if flag:
            # if not extract feature, send the request to server directly.
            request.body = json.dumps(data, ensure_ascii=False).encode('utf8')

    def run(self):
        """The entry to program start."""
        self.build()
        print('load model success')
        while True:
            uuid, request = self.input_queue.get()
            self.deal(uuid, request)


class TableHandler(tornado.web.RequestHandler):
    """The server of receive and deal request."""

    def initialize(self, input_queue: Queue,
                   futures_dict: Dict[str, concurrent.futures.Future]
                   ) -> None:
        """Initialize the request, called for each request.
        :param input_queue: The request queue.
        :param futures_dict: The future dict store the uuid of request and future.
        """
        self.input_queue = input_queue
        self.futures_dict = futures_dict

    def write(self, result: Union[str, 'Response']):
        if isinstance(result, Response):
            code, chunk = result.status_code, result.text
        else:
            code, chunk = 551, result
        self.set_status(code, reason=chunk)
        # chunk = json.dumps(result, ensure_ascii=False)
        # status = result.get('status', 0) or result.get('code', 0)
        # if status:
        #     self.set_status(status, reason=chunk)
        super(TableHandler, self).write(chunk)

    async def get(self):
        response = await self.deal()
        self.write(response)

    async def delete(self):
        response = await self.deal()
        self.write(response)

    async def put(self):
        response = await self.deal()
        self.write(response)

    async def post(self):
        response = await self.deal()
        self.write(response)

    async def deal(self):
        uuid = shortuuid.uuid()
        future = concurrent.futures.Future()
        self.futures_dict[uuid] = future
        self.input_queue.put((uuid, Request(self.request)))
        result = await wrap_future(future)
        return result


class Request(dict):
    def __init__(self, a: HTTPServerRequest):
        super(Request, self).__init__()
        self.method = a.method
        self.uri = a.uri
        self.headers = a.headers
        self.body = a.body

    def __str__(self):
        return f'{self.__class__.__name__}({self.__dict__})'

    def __repr__(self):
        return self.__str__()


class Response(dict):
    def __init__(self, a: requests.Response):
        super(Response, self).__init__()
        self.url = a.url
        self.text = a.text
        self.elapsed = a.elapsed
        self.status_code = a.status_code

    def __str__(self):
        return f'{self.__class__.__name__}({self.__dict__})'

    def __repr__(self):
        return self.__str__()


def set_result_thread(result_queue: Queue,
                      futures_dict: Dict[str, concurrent.futures.Future]
                      ) -> None:
    """Assign results asynchronously.
    :param result_queue: The queue of result which put by ``PackageProcess``.
    :param futures_dict: The dict which save future of each request.
    :return None.
    """
    while True:
        # uuid is the id of request.
        # flag is the status of request, True is success other is failed.
        # result is the result of request, if flag is false, result is the error info of request.
        uuid, flag, result = result_queue.get()
        if uuid in futures_dict:
            futures_dict[uuid].set_result(result)
            del futures_dict[uuid]


def run(port, url_queue, result_queue):
    tornado.options.parse_command_line()
    futures_dict = dict()
    t1 = threading.Thread(target=set_result_thread, args=(result_queue, futures_dict))
    t1.start()
    app = tornado.web.Application(
        [(f'/.*', TableHandler, dict(input_queue=url_queue, futures_dict=futures_dict))]
    )
    sockets = tornado.netutil.bind_sockets(port)
    server = tornado.httpserver.HTTPServer(app)
    server.add_sockets(sockets)
    tornado.ioloop.IOLoop.instance().start()


def install(model_name):
    if model_name == 'face_retrieval':
        util.install_package('tensorflow-gpu==1.15.2 mtcnn keras==2.2.4')
    elif model_name == 'image_retrieval':
        util.install_package('torchvision torch')
    elif model_name == 'text':
        text_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'text', 'requirements.txt')
        util.install_package(f'-r {text_path}')
    elif model_name == 'audio':
        raise NotImplementedError()
    else:
        raise Exception(f'{model_name} is not existed')


def main():
    define('model_name', type=str, default='face_retrieval', help='The model name you need')
    define('debug', type=bool, default=False, help='debug')
    tornado.options.parse_command_line()
    install(options.model_name)
    port = config.port
    gpus = config.gpus
    url_queue = Queue()
    result_queue = Queue()
    for gpu in gpus.split(','):
        package_process = PackageProcess(gpu, options.model_name, url_queue, result_queue, debug=options.debug)
        package_process.daemon = True
        package_process.start()

    run(port, url_queue, result_queue)


if __name__ == '__main__':
    main()