# gRPC imports
import grpc
from concurrent import futures

# classes generated by grpc
from .generated import eval_server_pb2
from .generated import eval_server_pb2_grpc

# Other imports
import numpy as np
import argparse
import pickle
from time import perf_counter
from os.path import join, isfile
from datetime import datetime
import os
import cmd
import time
import json

# multiprocessing
import multiprocessing as mp
from multiprocessing import shared_memory
from pycocotools.coco import COCO

# utils
from .utils import imread, gen_results, evaluate, mkdir2

FPS = 30

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data-root', type=str, required=True)
    parser.add_argument('--annot-path', type=str, required=True)
    parser.add_argument('--eval-config', type=str, required=False)
    parser.add_argument('--out-dir', type=str, required=True)
    parser.add_argument('--perf-factor', type=float, default=1)
    parser.add_argument('--overwrite', action='store_true', default=False)
    parser.add_argument('--log', type=str, default="server_log.txt")

    opts = parser.parse_args()
    return opts

class ImageServiceServicer(eval_server_pb2_grpc.ImageServiceServicer):

    def __init__(self, opts, db, current_sid, sequence_start_times_dict, log):
        self.opts = opts
        self.db = db
        self.seqs = self.db.dataset['sequences']
        self.seq_dirs = self.db.dataset['seq_dirs']
        self.fps = FPS

        img = self.db.imgs[0]
        self.img_width, self.img_height = img['width'], img['height']
        self.max_seq_len = 1000
        self.current_sid = current_sid
        self.sequence_start_times_dict = sequence_start_times_dict
        self.log = log

        # calculate the image size for shared memory allocation
        if 'max_seq_len' in self.db.dataset:
            self.max_seq_len = self.db.dataset['max_seq_len']
        else:
            seq_len = [len([img for img in self.db.imgs.values() if img['sid'] == sid])
                for sid in range(len(self.seqs))]
            self.max_seq_len = max(seq_len)

        self.shm = shared_memory.SharedMemory(create=True, size=(self.max_seq_len+2)*self.img_height*self.img_width*3)

    def close(self):
        self.shm.close()
        self.shm.unlink()

    def GetImageStream(self, request, context):
        
        n_total = 0
        for sid, seq in enumerate(self.seqs):
            if seq != request.value:
                continue
            print(datetime.now(), ": ", "Serving request - sid: " + str(sid) + " of " + str(len(self.seqs)-1), " name: ", request.value, file=self.log)
            frame_list = [img for img in self.db.imgs.values() if img['sid'] == sid]
    
            # load all frames in advance
            frames_local = np.ndarray((self.max_seq_len,) + (self.img_height, self.img_width, 3), dtype=np.uint8)
            frames = np.ndarray((self.max_seq_len,) + (self.img_height, self.img_width, 3), dtype=np.uint8, buffer=self.shm.buf)

            frames_permuted = np.random.permutation(len(frame_list))

            for i, img in enumerate(frame_list):
                if i % int(len(frame_list)/10) == 0:
                    print(datetime.now(), ": ", "Loaded frames (" + str(10*i/int(len(frame_list)/10)) + "%)", file=self.log)
                img_path = join(self.opts.data_root, self.seq_dirs[sid], img['name'])
                frames_local[frames_permuted[i]] = imread(img_path)
            frames[:] = frames_local[:]

            n_frame = len(frame_list)
            n_total += n_frame
            
            last_fidx = None
            
            t_total = (n_frame)/self.fps
            t_start = perf_counter()
            
            frames_processed = 0
            self.current_sid.value = sid
            self.sequence_start_times_dict[sid] = perf_counter()

            while 1:

                t1 = perf_counter()
                t_elapsed = t1 - t_start
                if t_elapsed >= t_total:
                    break

                # identify latest available frame
                fidx = int(np.floor(t_elapsed*self.fps))
                if fidx == last_fidx:
                    continue
                    
                frames_processed += 1
                last_fidx = fidx
                
                yield eval_server_pb2.Frame(fid = fidx, sid = sid, start_ptr=frames_permuted[fidx]*self.img_width*self.img_height*3, \
                                            end_ptr=(frames_permuted[fidx]+1)*self.img_width*self.img_height*3, timestamp=perf_counter())


            yield eval_server_pb2.Frame(end_marker=True, timestamp=perf_counter())
        print(datetime.now(), ": ","Finished streaming sequence ", request.value, file=self.log)

    def GetShm(self, request, context):
        return eval_server_pb2.String(value=self.shm.name)

class ResultServiceServicer(eval_server_pb2_grpc.ResultServiceServicer):

    def __init__(self, opts, db, current_sid, sequence_start_times_dict, log):
        self.opts = opts
        self.db = db
        self.seqs = self.db.dataset['sequences']
        self.sid_to_seq = {k:v for k,v in enumerate(self.seqs)}

        self.current_sid = current_sid
        self.sequence_start_times_dict = sequence_start_times_dict

        self.timestamps = []
        self.results_parsed = []
        self.input_fidx = []
        self.time_rcv = []
        self.log = log

        mkdir2(self.opts.out_dir)
    
    def reset_state(self):
        self.timestamps = []
        self.results_parsed = []
        self.input_fidx = []
        self.time_rcv = []

    def FinishSequence(self, result, context):
        seq_cur = self.sid_to_seq[self.current_sid.value]
        print(datetime.now(), ": ", "Writing result for sequence ", seq_cur, file=self.log)
        out_path = join(self.opts.out_dir, seq_cur + '.pkl')
        if self.opts.overwrite or not isfile(out_path):
            pickle.dump({
                'results_parsed': self.results_parsed,
                'timestamps': self.timestamps
            }, open(out_path, 'wb'))

        print(datetime.now(), ": ", "Mean time receive = ", np.mean(self.time_rcv), ", stdev = ", np.std(self.time_rcv), ", min = ", np.min(self.time_rcv), ", max = ", np.max(self.time_rcv), file=self.log)
        self.reset_state()
        return eval_server_pb2.Empty()

    def PutResultStream(self, result_iterator, context):
        for result in result_iterator:
            self.timestamps.append((perf_counter() - self.sequence_start_times_dict[self.current_sid.value])/self.opts.perf_factor)
            self.time_rcv.append(perf_counter() - result.timestamp)
            bboxes = np.ndarray((len(result.bboxes), 4), dtype=np.float32)
            bbox_scores = np.ndarray((len(result.bbox_scores),), dtype=np.float32)
            labels = np.ndarray((len(result.labels),), dtype=np.int64)
            for i in range(len(result.bboxes)):
                bboxes[i][0] = result.bboxes[i].x1
                bboxes[i][1] = result.bboxes[i].y1
                bboxes[i][2] = result.bboxes[i].x2
                bboxes[i][3] = result.bboxes[i].y2
                bbox_scores[i] = result.bbox_scores[i]
                labels[i] = result.labels[i]
            self.results_parsed.append((bboxes, bbox_scores, labels, None))
        return eval_server_pb2.Empty()

    def GenResults(self, request, context):
        gen_results(self.db, self.opts, request.value)
        return eval_server_pb2.Empty()

def result_server(opts, db, current_sid, sequence_start_times, res_rcv, log):
    print('Starting ResultService RPC Server. Listening on port ' + str(config['result_service_port']))

    result_server = grpc.server(futures.ThreadPoolExecutor(max_workers=12))
    result_servicer = ResultServiceServicer(opts, db, current_sid, sequence_start_times, log)
    eval_server_pb2_grpc.add_ResultServiceServicer_to_server(result_servicer, result_server)
    result_server.add_insecure_port('[::]:' + str(config['result_service_port']))
    result_server.start()
    msg = res_rcv.recv()
    if msg == 'terminate':
        result_server.stop(None)

class EvalShell(cmd.Cmd):
    intro = 'Welcome to the eval server.   Type help or ? to list commands.\n'
    prompt = '(eval_server) '

    def __init__(self, img_server, img_servicer, opts, db, res_send):
        super(EvalShell, self).__init__()
        self.db = db
        self.img_server = img_server
        self.img_servicer = img_servicer
        self.opts = opts
        self.res_send = res_send

    def do_close(self, arg):
        'Terminate evaluation server'
        self.img_servicer.close()
        self.img_server.stop(None)
        self.res_send.send('terminate')
        return True

    def do_gen_results(self, arg):
        'Generate and save results in coco-format to output directory. Usage: gen_results result_file.json'
        if not arg.endswith('.json') and not arg.endswith('.pkl'):
            print("Can only output to .json or .pkl file")
            return
        gen_results(self.db, self.opts, arg)
        print("Wrote results to: " + arg)

    def do_evaluate(self, arg):
        'Evaluate and display result metrics.'
        evaluate(self.db, self.opts.out_dir, arg, self.opts.overwrite)

    def do_log(self, arg):
        'View server log'
        os.system("cat " + self.opts.log)

if __name__ == '__main__':
    # create a gRPC server
    opts = parse_args()
    db = COCO(opts.annot_path)
    if opts.eval_config is not None:
        config = json.load(open(opts.eval_config, 'r'))
    else:
        config = {"image_service_port": 5007, "result_service_port": 5008, "loopback_ip": "127.0.0.1"}
        
    with open(opts.log, 'w+', 1) as log:
        # create dictionary to store current sequence and start times for each sequence stream
        current_sid = mp.Value('i', 0)
        sequence_start_times_dict = mp.Manager().dict()

        # start the image service RPC server
        server = grpc.server(futures.ThreadPoolExecutor(max_workers=12))
        servicer = ImageServiceServicer(opts, db, current_sid, sequence_start_times_dict, log)
        eval_server_pb2_grpc.add_ImageServiceServicer_to_server(servicer, server)
        server.add_insecure_port('[::]:' + str(config['image_service_port']))

        print('Starting ImageService RPC Server. Listening on port ' + str(config['image_service_port']))

        # start result service RPC server
        res_rcv, res_send = mp.Pipe(False)
        res_server = mp.Process(target=result_server, args=(opts, db, current_sid, sequence_start_times_dict, res_rcv, log))
        res_server.start()

        server.start()
        time.sleep(1)
        EvalShell(server, servicer, opts, db, res_send).cmdloop()
        server.wait_for_termination()




