# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function

import jupyter_client
import os
import sys
import threading
import time
from concurrent import futures

import grpc
import kernel_pb2
import kernel_pb2_grpc

is_py2 = sys.version[0] == '2'
if is_py2:
    import Queue as queue
else:
    import queue as queue


class KernelServer(kernel_pb2_grpc.JupyterKernelServicer):

    def __init__(self, server, kernel_name):
        self._status = kernel_pb2.STARTING
        self._server = server
        self._kernel_name = kernel_name
        # issue with execute_interactive and auto completion: https://github.com/jupyter/jupyter_client/issues/429
        # in all case because ipython does not support run and auto completion at the same time: https://github.com/jupyter/notebook/issues/3763
        # For now we will lock to ensure that there is no concurrent bug that can "hang" the kernel
        self._lock = threading.Lock()

    def start(self):
        print("starting...")
        sys.stdout.flush()
        self._km, self._kc = jupyter_client.manager.start_new_kernel(kernel_name=self._kernel_name)
        self._status = kernel_pb2.RUNNING

    def execute(self, request, context):
        # print("execute code:\n")
        # print(request.code.encode('utf-8'))
        sys.stdout.flush()
        stream_reply_queue = queue.Queue(maxsize = 30)
        payload_reply = []
        def _output_hook(msg):
            # print("msg: " + str(msg))
            msg_type = msg['header']['msg_type']
            content = msg['content']
            # print("******************")
            # print(msg)
            outStatus, outType, output = kernel_pb2.SUCCESS, None, None
            # prepare the reply
            if msg_type == 'stream':
                outType = kernel_pb2.TEXT
                output = content['text']
            elif msg_type in ('display_data', 'execute_result'):
                # print(content['data'])
                # The if-else order matters, can not be changed. Because ipython may provide multiple output.
                # TEXT is the last resort type.
                if 'text/html' in content['data']:
                    outType = kernel_pb2.HTML
                    output = content['data']['text/html']
                elif 'image/jpeg' in content['data']:
                    outType = kernel_pb2.JPEG
                    output = content['data']['image/jpeg']
                elif 'image/png' in content['data']:
                    outType = kernel_pb2.PNG
                    output = content['data']['image/png']
                elif 'application/javascript' in content['data']:
                    outType = kernel_pb2.HTML
                    output = '<script> ' + content['data']['application/javascript'] + ' </script>\n'
                elif 'application/vnd.holoviews_load.v0+json' in content['data']:
                    outType = kernel_pb2.HTML
                    output = '<script> ' + content['data']['application/vnd.holoviews_load.v0+json'] + ' </script>\n'
                elif 'text/plain' in content['data']:
                    outType = kernel_pb2.TEXT
                    output = content['data']['text/plain']
            elif msg_type == 'error':
                outStatus = kernel_pb2.ERROR
                outType = kernel_pb2.TEXT
                output = '\n'.join(content['traceback'])
            elif msg_type == 'clear_output':
                outType = kernel_pb2.CLEAR
                output = ""

            # send reply if we supported the output type
            if outType is not None:
                stream_reply_queue.put(
                    kernel_pb2.ExecuteResponse(status=outStatus,
                                                type=outType,
                                                output=output))
        def execute_worker():
            reply = self._kc.execute_interactive(request.code,
                                          output_hook=_output_hook,
                                          timeout=None)
            payload_reply.append(reply)

        t = threading.Thread(name="ConsumerThread", target=execute_worker)
        with self._lock:
            t.start()
            # We want to wait the end of the execution (and queue empty).
            # In our case when the thread is not alive -> it means that the execution is complete
            # However we also ensure that the kernel is alive because in case of OOM or other errors
            # Execution might be stuck there: (might open issue on jupyter client)
            # https://github.com/jupyter/jupyter_client/blob/master/jupyter_client/blocking/client.py#L323
            while (t.is_alive() and self.isKernelAlive()) or not stream_reply_queue.empty():
                # Sleeping time to time to reduce cpu usage.
                # At worst it will bring a 0.05 delay for bunch of messages.
                # Overall it will improve performance.
                time.sleep(0.05)
                while not stream_reply_queue.empty():
                    yield stream_reply_queue.get()

            # if kernel is not alive or thread is still alive, it means that we face an issue.
            if not self.isKernelAlive() or t.is_alive():
                yield kernel_pb2.ExecuteResponse(status=kernel_pb2.ERROR,
                                                  type=kernel_pb2.TEXT,
                                                  output="Ipython kernel has been stopped. Please check logs. It might be because of an out of memory issue.")
        if payload_reply:
            result = []
            if 'payload' in payload_reply[0]['content']:
                for payload in payload_reply[0]['content']['payload']:
                    if payload['data']['text/plain']:
                        result.append(payload['data']['text/plain'])
            if result:
                yield kernel_pb2.ExecuteResponse(status=kernel_pb2.SUCCESS,
                                                  type=kernel_pb2.TEXT,
                                                  output='\n'.join(result))

    def cancel(self, request, context):
        self._km.interrupt_kernel()
        return kernel_pb2.CancelResponse()

    def complete(self, request, context):
        with self._lock:
            reply = self._kc.complete(request.code, request.cursor, reply=True, timeout=None)
        return kernel_pb2.CompletionResponse(matches=reply['content']['matches'])

    def status(self, request, context):
        return kernel_pb2.StatusResponse(status = self._status)

    def isKernelAlive(self):
        return self._km.is_alive()

    def terminate(self):
        self._km.shutdown_kernel()

    def stop(self, request, context):
        self.terminate()
        return kernel_pb2.StopResponse()


def serve(kernel_name, port):
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    kernel = KernelServer(server, kernel_name)
    kernel_pb2_grpc.add_JupyterKernelServicer_to_server(kernel, server)
    server.add_insecure_port('[::]:' + port)
    server.start()
    kernel.start()
    try:
        while kernel.isKernelAlive():
            time.sleep(5)
    except KeyboardInterrupt:
        print("interrupted")
    finally:
        print("shutdown")
        # we let 2 sc for all request to be complete
        server.stop(2)
        kernel.terminate()
        os._exit(0)

if __name__ == '__main__':
    serve(sys.argv[1], sys.argv[2])
