import os
import threading
import traceback

import grpc
from concurrent import futures
import config_pb2_grpc
import config_pb2

from app import app_run, get_parameters

import logging

import sys


def my_except_hook(exctype, value, traceback):
    if exctype == KeyboardInterrupt:
        print("Handler code goes here")
    else:
        sys.__excepthook__(exctype, value, traceback)


sys.excepthook = my_except_hook

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

port = 8061
shared_folder = os.getenv("SHARED_FOLDER_PATH")
print(f'Shared folder: {shared_folder}')


class TrainingConfiguration(config_pb2_grpc.TrainingConfigurationServicer):

    def __init__(self):
        super().__init__()
        self.send_data = True
        self.start_count = 0

    def startTraining(self, request, context):
        response = config_pb2.TrainingConfig()
        print("Entered start Training")
        if self.start_count > 0:
            if self.start_count == 1:
                logger.debug("Training already completed.")
            context.set_code(grpc.StatusCode.NOT_FOUND)
            context.set_details("Training completed.")
        else:
            try:
                logger.debug("Start the training...")
                params = get_parameters()
                response.training_data_filename = params.training_data_filename.data.filename
                response.validation_data_filename = params.validation_data_filename.data.filename
                response.test_data_filename = params.test_data_filename.data.filename

                for embedding in params.embeddings:
                    response.embeddings.extend([config_pb2.Embedding(embedding=embedding)])

                response.epochs = params.epochs
                response.learning_rate = params.learning_rate
                response.batch_size = params.batch_size
                response.model_filename = params.model_filename
            except Exception as e:
                logging.error(traceback.format_exc())

        self.start_count += 1

        if not self.send_data:
            context.set_code(grpc.StatusCode.NOT_FOUND)
            context.set_details("Data has been processed")

        self.send_data = not self.send_data

        return response


server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
config_pb2_grpc.add_TrainingConfigurationServicer_to_server(TrainingConfiguration(), server)
logger.debug("Starting grpc server. Listening on port : " + str(port))
server.add_insecure_port("[::]:{}".format(port))
server.start()
logger.debug("Start ner-training-config web-ui")
threading.Thread(target=app_run()).start()
server.wait_for_termination()