"""
eui means easy UI. It is a fast and simple micro-framework for small browser-based applications.
Python and JS communicate through websocket protocol.
"""

import base64
import hashlib
import json
import logging
import mimetypes
import os
import queue
import random
import socket
import struct
import threading
import time
from concurrent import futures

log = logging.getLogger('eui')

# eui.js code template
_JS_TEMPLATE = '''/**
* This file is automatically generated every time eui starts. Please do not modify it manually.
 */
window.eui = {};
eui.retryTimes = 10;

eui.init = function () {
   eui.ws = new WebSocket("ws://localhost:%s");
   eui.ws.onopen = function () {
      console.log('connect to eui server!');
      eui.retryTimes = 10;
   };

   eui.ws.onmessage = function (evt) {
      //heartbeat
      if (evt.data == 'pong') { return }
      try {
         var data = JSON.parse(evt.data);
         var handler = eval(data['handler']);
         if (typeof handler != 'function') {
            eui.onerror(data['handler'] + ' is not a valid function');
            return;
         }
         if (data['args']){
            handler(...data['args']);
         } else {
            handler();
         }
      } catch (e) {
         eui.onerror(e);
      }
   };

   eui.ws.onclose = function () {};

   eui.ws.onerror = function () {};

}

/**
* It is used to execute Python handler. 
* 
* @param {string} handler Python handler
* @param {...} ...args args for Python handler. If there are multiple args, separate them with commas.
*/
eui.py = function (handler, ...args) {
  if (eui.ws.readyState != 1) {
     setTimeout(eui.py, 100, handler, ...args);
  } else {
     eui.ws.send(JSON.stringify({ 'handler': handler, 'args': args }));
  }
};

/**
* Default error handler. It will be executed when an exception occurs and can be overridden if necessary.
* 
* @param {string} e error message
*/
eui.onerror = function (e) {
  alert(e);
  console.error(e);
};

/**
* exit app
* 
* @param {code} error code
*/
eui.exit = function (code) {
  eui.py('exit', 0);
};

if ('WebSocket' in window) {
   eui.init();
   //send heartbeat
   setInterval(() => {
       eui.ws.send('ping');
    }, 6000);

    //check connection
    setInterval(() => {
       if (eui.retryTimes <= 0) {
          window.close();
       } else if (eui.ws && eui.ws.readyState > 1) {
          console.log('connection closed, retry init...');
          eui.retryTimes -= 1;
          eui.init();
       }
    }, 3000);
} else {
   alert("your browser don't support WebSocket");
}

'''


class Eui:
    def __init__(self, handlers, callback=None, port=None, static_dir='./static', max_workers=10, timeout=30,
                 log_file='eui.log', log_level='DEBUG', max_message_size=1024 * 1024):
        """
        init eui

        :param handlers: python function for js call
        :param callback: the function after app startup to run
        :param port: port, if port is None, port will be a random int value
        :param static_dir: dir for output eui.js
        :param max_workers: max workers to execute handlers
        :param timeout: max seconds to wait for heartbeat
        :param log_file: log file path
        :param log_level: log level, 'CRITICAL', 'FATAL', 'ERROR', 'WARN', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'
        :return:
        """
        self.__send_queue = queue.Queue()
        self.handlers = handlers
        self.callback = callback
        if not port:
            port = random.randint(5000, 60000)
        self.port = port
        self.__workers = futures.ThreadPoolExecutor(max_workers)
        self.__max_message_size = max_message_size
        self.__timeout = timeout
        self.__ttl = self.__timeout
        self.__ws_connection = None
        self.__init_log(log_file, log_level)
        self.__init_js(static_dir)

    @staticmethod
    def __init_log(file, level):
        """
        init log config
        :param file: log file
        :param level: log level
        :return:
        """
        log.setLevel(level)
        formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(pathname)s [line:%(lineno)d] %(message)s",
                                      '%Y-%m-%d %H:%M:%S')
        console_log_handler = logging.StreamHandler()
        console_log_handler.setFormatter(formatter)
        file_log_handler = logging.FileHandler(file, mode='w')
        file_log_handler.setFormatter(formatter)

        log.addHandler(console_log_handler)
        log.addHandler(file_log_handler)

    @staticmethod
    def __get_headers(payload):
        """
        get headers from data
        :param payload: received data
        :return: headers
        """
        headers = {}
        if not payload:
            return headers
        header_str, body = str(payload, encoding="utf-8").split("\r\n\r\n", 1)
        header_list = header_str.split("\r\n")
        headers['method'], headers['uri'], headers['protocol'] = header_list[0].split(' ')
        for row in header_list[1:]:
            key, value = row.split(":", 1)
            headers[str(key).lower()] = value.strip()

        return headers

    @staticmethod
    def __parse_message(payload: bytes) -> str:
        """
        parse payload message

        :param payload: bytes message
        :return: parsed string message
        """
        if not payload:
            return ''

        try:
            payload_len = payload[1] & 127
            if payload_len == 126:
                mask = payload[4:8]
                decoded = payload[8:]

            elif payload_len == 127:
                mask = payload[10:14]
                decoded = payload[14:]
            else:
                mask = payload[2:6]
                decoded = payload[6:]

            byte_list = bytearray()
            for i, b in enumerate(decoded):
                chunk = b ^ mask[i % 4]
                byte_list.append(chunk)

            if byte_list == b'\x03\xe9':
                return ''

            return str(byte_list, encoding='utf-8')
        except Exception as e:
            log.warning('parse payload error: %s', e)
            return ''

    def __init_js(self, static_dir):
        """
        generate js file

        :param static_dir: dir for eui.js
        :return:
        """
        os.makedirs(static_dir, exist_ok=True)
        path = (static_dir if static_dir.endswith('/') else static_dir + '/') + 'eui.js'
        with open(path, 'w', encoding='utf-8') as f:
            f.write(_JS_TEMPLATE % self.port)
        log.debug('init js path: %s', path)

    def __do_work(self, message):
        """
        message handler
        :param message: request message
        :return:
        """
        try:
            message = self.__parse_message(message)
            # connection close
            if not message:
                time.sleep(0.1)
                return
            log.debug('receive message: %s', message)
            # if is heart beat message, reset ttl
            if message == 'ping':
                self.__ttl = self.__timeout
                self.__send_message(b'pong')
                log.debug('send message: pong')
                return
            # parse json data
            message = json.loads(message)
            handler_name = message.get('handler', '')
            if handler_name not in self.handlers:
                self.error("handler '%s' not found, please check Python handlers config" % handler_name)
                return
            handler = self.handlers[handler_name]
            args = message.get('args', None)
            if args:
                handler(*args)
            else:
                handler()
        except Exception as e:
            self.error(str(e), e)

    def __startup_receiver(self):
        """
        startup a worker to receive message
        """

        def run():
            while True:
                if not self.__ws_connection:
                    time.sleep(0.1)
                    continue
                try:
                    message = self.__ws_connection.recv(self.__max_message_size)
                    self.__workers.submit(self.__do_work, message)
                except ConnectionAbortedError:
                    log.warning('connection aborted')
                    time.sleep(1)
                except Exception as e:
                    log.warning('receive message error', exc_info=e)

        # send message worker is a daemon thread
        threading.Thread(target=run, daemon=True).start()

    def __startup_sender(self):
        """
        startup a worker to send message

        """

        def run():
            while True:
                if not self.__ws_connection:
                    time.sleep(0.1)
                    continue
                try:
                    message = self.__send_queue.get().encode('utf-8')
                    self.__send_message(message)
                except Exception as e:
                    log.error('eui send message error: %s', str(e), exc_info=e)

        threading.Thread(target=run, daemon=True).start()

    def __send_message(self, message: bytes):
        """send message to js"""
        if not self.__ws_connection:
            log.warning('send message error, no usable connection')
            return

        try:
            token = b"\x81"
            length = len(message)
            if length < 126:
                token += struct.pack("B", length)
            elif length <= 0xFFFF:
                token += struct.pack("!BH", 126, length)
            else:
                token += struct.pack("!BQ", 127, length)

            msg = token + message
            self.__ws_connection.sendall(msg)
        except ConnectionAbortedError:
            pass
        except Exception as e:
            log.error('send message error: %s', str(e), exc_info=e)

    @staticmethod
    def __startup_callback(fn):
        """
        startup callback function, when eui startup, fn will be execute
        :param fn: callback function
        :return:
        """
        if fn:
            threading.Thread(target=fn, daemon=True).start()

    def __startup_timer(self):
        """
        startup a worker to check timeout event
        """
        if not self.__timeout:
            return

        def run():
            while True:
                if self.__ttl <= 0:
                    log.info('wait for connect timeout, app exit!')
                    self.exit()
                self.__ttl -= 1
                time.sleep(1)

        threading.Thread(target=run, daemon=True).start()

    def __new_websocket_connection(self, headers, connection):
        """
        create new websocket connection
        :param headers: request headers
        :param connection: websocket connection
        :return:
        """
        key = ''
        if headers.get('sec-websocket-key'):
            key = headers['sec-websocket-key'] + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
        ac = base64.b64encode(hashlib.sha1(key.encode('utf-8')).digest())
        response = "HTTP/1.1 101 Switching Protocols\r\n" \
                   "Upgrade: websocket\r\n" \
                   "Connection: Upgrade\r\n" \
                   "Sec-WebSocket-Accept: %s\r\n" \
                   "WebSocket-Location: ws://%s\r\n\r\n" % (ac.decode('utf-8'), headers.get("host"))
        connection.sendall(bytes(response, encoding="utf-8"))

        # close old connection
        if self.__ws_connection:
            try:
                self.__ws_connection.close()
            except Exception as e:
                log.warning('connection close error: %s' % str(e))

        self.__ws_connection = connection

    @staticmethod
    def __new_http_connection(headers, connection):
        """
        create new http connection
        :param headers: request headers
        :param connection: http connection
        :return:
        """
        path = '.' + headers.get('uri', '')
        if os.path.exists(path) and os.path.isfile(path):
            mimetype, _ = mimetypes.guess_type(path)
            with open(path, 'rb') as f:
                data = f.read()
                response = "HTTP/1.1 200 OK\r\n" \
                           "Content-Type: %s\r\n" \
                           "Content-Length: %s\r\n\r\n" % (mimetype, len(data))
                connection.sendall(bytes(response, encoding='utf-8') + data)
        else:
            log.error("can't read file '%s'", path)
            connection.sendall(b"HTTP/1.1 404 Not Found")

        connection.close()

    def __create_connection(self, connection):
        """
        create websocket or http connection
        :param connection: socket connection
        :return:
        """
        message = connection.recv(self.__max_message_size)
        headers = self.__get_headers(message)
        if not headers:
            connection.sendall(b"HTTP/1.1 200 OK")
            connection.close()
            return

        # websocket connection
        if headers.get('upgrade', '').lower() == 'websocket':
            self.__new_websocket_connection(headers, connection)
        else:
            # http connection
            self.__new_http_connection(headers, connection)

    def error(self, message, err=None):
        """
        send error message to js
        :param message: error message
        :param err: error object
        :return:
        """
        log.error(message, exc_info=err)
        self.js('eui.onerror', message)

    def js(self, handler, *args):
        """
        call js function

        :param handler: js function
        :param args: js function args
        :return:
        """
        self.__send_queue.put(json.dumps({'handler': handler, 'args': args}, ensure_ascii=True))

    def exit(self):
        """
        exit app
        :return:
        """
        self.__send_message(json.dumps({'handler': 'window.close'}).encode('utf-8'))
        log.info('app exit!')
        os.abort()

    def run(self):
        """
        run app

        """
        # init socket
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind(('127.0.0.1', self.port))
        sock.listen(6)
        self.__startup_receiver()
        self.__startup_sender()
        self.__startup_callback(self.callback)
        self.__startup_timer()
        print('******************** eui start up at port %s ********************' % self.port)

        # accept connection
        while True:
            try:
                connection, address = sock.accept()
                self.__workers.submit(self.__create_connection, connection)
            except Exception as e:
                log.error('create connection error', exc_info=e)


_app: Eui = None


def start(handlers=None, callback=None, port=None, static_dir='./static', max_workers=10, timeout=30,
          log_file='eui.log', log_level='DEBUG', max_message_size=1024 * 1024):
    """
    startup eui

    :param handlers: python function for js call
    :param callback: the function after app startup to run
    :param port: port, if port is None, port will be a random int value
    :param static_dir: dir for output eui.js
    :param max_workers: max workers to execute handlers
    :param timeout: max seconds to wait for heartbeat
    :param log_file: log file path
    :param log_level: log level, 'CRITICAL', 'FATAL', 'ERROR', 'WARN', 'WARNING', 'INFO', 'DEBUG', 'NOTSET'
    :param max_message_size: max receive message size
    :return:
    """

    args = locals()
    global _app
    _app = Eui(handlers, callback, port, static_dir, max_workers, timeout, log_file, log_level, max_message_size)
    _app.run()


def js(handler, *args):
    """
    call js function

    :param handler: js function
    :param args: js function args
    :return:
    """
    _app.js(handler, *args)


def error(message, err):
    """
    send error message to js
    :param message: error message
    :param err: error object
    :return:
    """
    _app.error(message, err)


def get_port():
    """
    get app port
    :return: port
    """
    return _app.port


def exit():
    """
    quit app
    :return:
    """
    _app.exit()
