import argparse
import logging
import time
import os
import sys
import random
import math
import traceback
import uuid
from functools import reduce
import inspect
import threading
import copy
from imjoySocketIO_client import SocketIO, LoggingNamespace, find_callback
from imjoyUtils import debounce, setInterval, dotdict, ReferenceStore

if sys.version_info >= (3, 0):
    import asyncio
    import janus
    from imjoyUtils3 import task_worker, FuturePromise
    PYTHON3 = True
else:
    from imjoyUtils import task_worker, Promise
    PYTHON3 = False

try:
    import queue
except ImportError:
    import Queue as queue

logging.basicConfig(stream=sys.stdout)
logger = logging.getLogger('plugin')
logger.setLevel(logging.INFO)
# import logging
# logging.basicConfig(level=logging.DEBUG)
ARRAY_CHUNK = 1000000

if '' not in sys.path:
    sys.path.insert(0, '')

imjoy_path = os.path.dirname(os.path.normpath(__file__))
if imjoy_path not in sys.path:
    sys.path.insert(0, imjoy_path)

def kill(proc_pid):
    import psutil
    process = psutil.Process(proc_pid)
    for proc in process.children(recursive=True):
        proc.kill()
    process.kill()

def ndarray(typedArray, shape, dtype):
    _dtype = type(typedArray)
    if dtype and dtype != _dtype:
        raise Exception("dtype doesn't match the type of the array: "+_dtype+' != '+dtype)
    shape = shape or (len(typedArray), )
    return {"__jailed_type__": 'ndarray', "__value__" : typedArray, "__shape__": shape, "__dtype__": _dtype}

api_utils = dotdict(ndarray=ndarray, kill=kill, debounce=debounce, setInterval=setInterval)

class PluginConnection():
    def __init__(self, pid, secret, server, queue=None, loop=None, worker=None, namespace='/', work_dir=None, daemon=False, api=None):
        if work_dir is None or work_dir == '' or work_dir == '.':
            self.work_dir = os.getcwd()
        else:
            self.work_dir = work_dir
            if not os.path.exists(self.work_dir):
                os.makedirs(self.work_dir)
            os.chdir(self.work_dir)
        socketIO = SocketIO(server, Namespace=LoggingNamespace)
        self.socketIO = socketIO
        self._init = False
        self.secret = secret
        self.id = pid
        self.daemon = daemon

        def emit(msg):
            socketIO.emit('from_plugin_'+ secret, msg)
        self.emit = emit

        self._local = {}
        _remote = dotdict()
        self._setLocalAPI(_remote)
        self._interface = {}
        self._plugin_interfaces = {}
        self._remote_set = False
        self._store = ReferenceStore()
        self._executed = False
        self.queue = queue
        self.loop = loop

        self._init = False
        sys.stdout.flush()
        socketIO.on('to_plugin_'+secret, self.sio_plugin_message)
        self.emit({"type": "initialized", "dedicatedThread": True})
        print('Plugin "{}" Initialized.'.format(pid))
        def on_disconnect():
            if not self.daemon:
                self.exit(1)
        socketIO.on('disconnect', on_disconnect)
        self.abort = threading.Event()
        self.worker = worker


    def wait_forever(self):
        if PYTHON3:
            self.sync_q = self.queue.sync_q
            fut = self.loop.run_in_executor(None, self.socketIO.wait)
            t = [self.worker(self, self.queue.async_q, logger, self.abort) for i in range(10)]
            self.loop.run_until_complete(asyncio.gather(*t))
            self.loop.run_until_complete(fut)
        else:
            self.sync_q = queue.Queue()
            t = threading.Thread(target=self.socketIO.wait)
            t.daemon = True
            t.start()
            self.worker(self, self.sync_q, logger, self.abort)

    def default_exit(self):
        logger.info('terminating plugin: ' + self.id)
        self.abort.set()
        os._exit(0)

    def exit(self, code):
        if 'exit' in self._interface:
            try:
                self._interface['exit']()
            except Exception as e:
                logger.error('Error when exiting: %s', e)
                sys.exit(1)
            else:
                logger.info('terminating plugin')
                sys.exit(code)
        else:
            sys.exit(0)

    def _encode(self, aObject):
        if aObject is None:
            return aObject
        if type(aObject) is tuple:
            aObject = list(aObject)
        isarray = type(aObject) is list
        bObject = [] if isarray else {}
        #skip if already encoded
        if type(aObject) is dict and '__jailed_type__' in aObject and '__value__' in aObject:
            return aObject

        # encode interfaces
        if type(aObject) is dict and '__id__' in aObject and '__jailed_type__' in aObject and aObject['__jailed_type__'] == 'plugin_api':
            encoded_interface = {}
            for k in aObject.keys():
                v = aObject[k]
                if callable(v):
                    bObject[k] = {'__jailed_type__': 'plugin_interface', '__plugin_id__':aObject['__id__'], '__value__' : k, 'num': None}
                    encoded_interface[k] = v
            self._plugin_interfaces[aObject['__id__']] = encoded_interface
            return bObject

        keys = range(len(aObject)) if isarray else aObject.keys()
        for k in keys:
            v = aObject[k]
            value = None
            try:
                basestring
            except NameError:
                basestring = str
            if callable(v):
                interfaceFuncName = None
                for name in self._interface:
                    if self._interface[name] == v:
                        interfaceFuncName = name
                        break
                if interfaceFuncName is None:
                    cid = self._store.put(v)
                    vObj = {'__jailed_type__': 'callback', '__value__' : 'f', 'num': cid}
                else:
                    vObj = {'__jailed_type__': 'interface', '__value__' : interfaceFuncName}

          # // send objects supported by structure clone algorithm
          # // https://developer.mozilla.org/en-US/docs/Web/API/Web_Workers_API/Structured_clone_algorithm
            #if(v !== Object(v) || v instanceof Boolean || v instanceof String || v instanceof Date || v instanceof RegExp || v instanceof Blob || v instanceof File || v instanceof FileList || v instanceof ArrayBuffer || v instanceof ArrayBufferView || v instanceof ImageData){
            elif 'np' in self._local and isinstance(v, (self._local['np'].ndarray, self._local['np'].generic)):
                vb = bytearray(v.tobytes())
                if len(vb)>ARRAY_CHUNK:
                    vl = int(math.ceil(1.0*len(vb)/ARRAY_CHUNK))
                    v_bytes = []
                    for i in range(vl):
                        v_bytes.append(vb[i*ARRAY_CHUNK:(i+1)*ARRAY_CHUNK])
                else:
                    v_bytes = vb
                vObj = {'__jailed_type__': 'ndarray', '__value__' : v_bytes, '__shape__': v.shape, '__dtype__': str(v.dtype)}
            elif type(v) is dict or type(v) is list:
                vObj = self._encode(v)
            elif not isinstance(v, basestring) and type(v) is bytes:
                vObj = v.decode() # covert python3 bytes to str
            elif isinstance(v, Exception):
                vObj = {'__jailed_type__': 'error', '__value__' : str(v)}
            else:
                vObj = {'__jailed_type__': 'argument', '__value__' : v}

            if isarray:
                bObject.append(vObj)
            else:
                bObject[k] = vObj

        return bObject

    def _decode(self, aObject, callbackId, withPromise):
        if aObject is None:
            return aObject
        if '__jailed_type__' in aObject and '__value__' in aObject:
            if aObject['__jailed_type__'] == 'callback':
                bObject = self._genRemoteCallback(callbackId, aObject['num'], withPromise)
            elif aObject['__jailed_type__'] == 'interface':
                name = aObject['__value__']
                if name in self._remote:
                    bObject = self._remote[name]
                else:
                    bObject = self._genRemoteMethod(name)
            elif aObject['__jailed_type__'] == 'plugin_interface':
                bObject = self._genRemoteMethod(aObject['__value__'], aObject['__plugin_id__'])
            elif aObject['__jailed_type__'] == 'ndarray':
                # create build array/tensor if used in the plugin
                try:
                    np = self._local['np']
                    if isinstance(aObject['__value__'], bytearray):
                        aObject['__value__'] = aObject['__value__']
                    elif isinstance(aObject['__value__'], list) or isinstance(aObject['__value__'], tuple):
                        aObject['__value__'] = reduce((lambda x, y: x + y), aObject['__value__'])
                    else:
                        raise Exception('Unsupported data type: ', type(aObject['__value__']), aObject['__value__'])
                    bObject = np.frombuffer(aObject['__value__'], dtype=aObject['__dtype__']).reshape(tuple(aObject['__shape__']))
                except Exception as e:
                    logger.debug('Error in converting: %s', e)
                    bObject = aObject
                    raise e
            elif aObject['__jailed_type__'] == 'error':
                bObject = Exception(aObject['__value__'])
            elif aObject['__jailed_type__'] == 'argument':
                bObject = aObject['__value__']
            else:
                bObject = aObject['__value__']
            return bObject
        else:
            if type(aObject) is tuple:
                aObject = list(aObject)
            isarray = type(aObject) is list
            bObject =  [] if isarray else dotdict()
            keys = range(len(aObject)) if isarray else aObject.keys()
            for k in keys:
                if isarray or k in aObject:
                    v = aObject[k]
                    if isinstance(v, dict)or type(v) is list:
                        if isarray:
                            bObject.append(self._decode(v, callbackId, withPromise))
                        else:
                            bObject[k] = self._decode(v, callbackId, withPromise)
            return bObject

    def _wrap(self, args):
        wrapped = self._encode(args)
        result = {'args': wrapped}
        return result

    def _unwrap(self, args, withPromise):
        called = False
        if "callbackId" not in args:
            args["callbackId"] = None
        # wraps each callback so that the only one could be called
        result = self._decode(args["args"], args["callbackId"], withPromise)
        return result

    def setInterface(self, api):
        if inspect.isclass(type(api)):
            api = {a:getattr(api, a) for a in dir(api) if not a.startswith('_')}
        if 'exit' in api:
            ext = api['exit']
            def exit_wrapper():
                try:
                    ext()
                finally:
                    self.default_exit()
            api['exit'] = exit_wrapper
        else:
            api['exit'] = self.default_exit
        self._interface = api
        self._sendInterface()

    def _sendInterface(self):
        names = []
        for name in self._interface:
            if callable(self._interface[name]):
                names.append({"name":name, "data": None})
            else:
                data = self._interface[name]
                if data is not None and type(data) is dict:
                    data2 = {}
                    for k in data:
                        if callable(data[k]):
                            data2[k] = "**@@FUNCTION@@**:"+k
                        else:
                            data2[k] = data[k]
                    names.append({"name":name, "data": data2})
                elif type(data) in [str, int, float, bool]:
                  names.append({"name":name, "data": data})
        self.emit({'type':'setInterface', 'api': names})

    def _genRemoteMethod(self, name, plugin_id=None):
        def remoteMethod(*arguments, **kwargs):
            # wrap keywords to a dictionary and pass to the first argument
            if len(arguments) == 0 and len(kwargs) > 0:
                arguments = [kwargs]
            def p(resolve, reject):
                resolve.__jailed_pairs__ = reject
                reject.__jailed_pairs__ = resolve
                call_func = {
                    'type': 'method',
                    'name': name,
                    'pid': plugin_id,
                    'args': self._wrap(arguments),
                    'promise': self._wrap([resolve, reject])
                }
                self.emit(call_func)
            if PYTHON3:
                return FuturePromise(p, self.loop)
            else:
                return Promise(p)
        remoteMethod.__remote_method = True
        return remoteMethod

    def _genRemoteCallback(self, id, argNum, withPromise):
        if withPromise:
            def remoteCallback(*arguments, **kwargs):
                # wrap keywords to a dictionary and pass to the first argument
                if len(arguments) == 0 and len(kwargs) > 0:
                    arguments = [kwargs]
                def p(resolve, reject):
                    resolve.__jailed_pairs__ = reject
                    reject.__jailed_pairs__ = resolve
                    self.emit({
                        'type' : 'callback',
                        'id'   : id,
                        'num'  : argNum,
                        # 'pid'  : self.id,
                        'args' : self._wrap(arguments),
                        'promise': self._wrap([resolve, reject])
                    })
                if PYTHON3:
                    return FuturePromise(p, self.loop)
                else:
                    return Promise(p)

        else:
            def remoteCallback(*arguments, **kwargs):
                # wrap keywords to a dictionary and pass to the first argument
                if len(arguments) == 0 and len(kwargs) > 0:
                    arguments = [kwargs]
                ret = self.emit({
                    'type' : 'callback',
                    'id'   : id,
                    'num'  : argNum,
                    # 'pid'  : self.id,
                    'args' : self._wrap(arguments)
                })
                return ret
        return remoteCallback

    def _setRemote(self, api):
        _remote = dotdict()
        for i in range(len(api)):
            if type(api[i]) is dict and "name" in api[i]:
                name = api[i]["name"]
                data = api[i].get("data", None)
                if data is not None:
                    if type(data) is dict:
                        data2 = dotdict()
                        for key in data:
                            if key in data:
                                if data[key] == "**@@FUNCTION@@**:"+key:
                                    data2[key] = self._genRemoteMethod(name+'.'+key)
                                else:
                                    data2[key] = data[key]
                        _remote[name] = data2
                    else:
                        _remote[name] = data
                else:
                    _remote[name] = self._genRemoteMethod(name)

        self._setLocalAPI(_remote)
        return _remote

    def _setLocalAPI(self, _remote):
        _remote["export"] = self.setInterface
        _remote["utils"] = api_utils
        _remote["WORK_DIR"] = self.work_dir
        self._local["api"] = _remote

    def sio_plugin_message(self, *args):
        data = args[0]
        if data['type']== 'import':
            self.emit({'type':'importSuccess', 'url': data['url']})
        elif data['type']== 'disconnect':
            self.abort.set()
            callback, args = find_callback(args)
            try:
                if 'exit' in self._interface and callable(self._interface['exit']):
                    self._interface['exit']()
            except Exception as e:
                logger.error('Error when exiting: %s', e)
            if callback:
                callback(*args)
        elif data['type'] == 'execute':
            if not self._executed:
                self.sync_q.put(data)
            else:
                logger.debug('skip execution.')
                self.emit({'type':'executeSuccess'})
        elif data['type'] == 'message':
            d = data['data']
            self.sync_q.put(d)
            logger.debug('added task to the queue')
        sys.stdout.flush()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--id', type=str, required=True, help='plugin id')
    parser.add_argument('--secret', type=str, required=True, help='plugin secret')
    parser.add_argument('--namespace', type=str, default='/', help='socketio namespace')
    parser.add_argument('--work_dir', type=str, default='.', help='plugin working directory')
    parser.add_argument('--server', type=str, default='http://127.0.0.1:9527', help='socketio server')
    parser.add_argument('--daemon', action="store_true", help='daemon mode')
    parser.add_argument('--debug', action="store_true", help='debug mode')

    opt = parser.parse_args()
    if opt.debug:
        logger.setLevel(logging.DEBUG)

    if PYTHON3:
        loop = asyncio.get_event_loop()
        q = janus.Queue(loop=loop)
    else:
        loop = None
        q = None

    pc = PluginConnection(opt.id, opt.secret, server=opt.server, work_dir=opt.work_dir, queue=q, loop=loop, worker=task_worker)
    pc.wait_forever()
