"""
Python client based on https://pypi.org/project/websocket-client/

@author hornetfighter515
@author Todorov-Lachezar 
"""

import websocket
import _thread
import time
import ssl
import yaml
import json
import rsa
import random # for message id generation
import cipher

ext = False
authed = False
clis = []
messages = {}

def authed_ext_send(ws, msg):
    msg = rsa.encrypt(msg.encode(), clis[0]).decode('latin')
    message_id = random.random()
    if len(clis) <= 1: # if there are no IoTs
        # send the message once-encrypted
        message = cipher.gen_comm_msg(0, message_id, msg)
        ws.send(message)
        return
    # otherwise, slice and twice encrypt it
    twice_es = []
    modifier = len(msg)/(len(clis)-1)
    for i in range(1,len(clis)):
        start = int( (i-1) * modifier)
        end = int(i*modifier)
        # makes a slice of the encrypted message, and indicates the intended 
        # recipient
        contents = rsa.encrypt(msg[start:end].encode('latin'), 
                clis[i]).decode('latin')
        message = cipher.gen_comm_msg(i, message_id, contents)
        twice_es.append(message)
    for m in twice_es:
        ws.send(m)


def authed_int_send(ws, msg):
    to = int(input(f"Who's it to? 0-{len(clis)}>"))
    e_msg = rsa.encrypt(msg.encode(), clis[int(to)]).decode('latin')
    message_id = random.random()
    msg = cipher.gen_comm_msg(to, message_id, e_msg)
    ws.send(msg)


def init_send(ws):
    message = cipher.gen_pub_key_msg(pub['n'], pub['e'])
    ws.send(message)


def send(ws, msg):
    try:
        if ext:
            authed_ext_send(ws, msg)
        else:
            authed_int_send(ws, msg)
    except Exception as e:
        print(f'Sending message {msg} failed: {e}')


def on_message(ws, message):
    msg = json.loads(message)
    global authed
    if msg['type'] == 'pub_key':
        authed = True
        if ext:
            k = rsa.PublicKey(msg['available_cli']['n'], 
                    msg['available_cli']['e'])
            if len(clis) == 0:
                clis.append(k)
            else:
                clis[0] = k
            for i,p in msg['available_iots'].items():
                k = rsa.PublicKey(p['n'], p['e'])
                clis.append(k)
        else:
            k = rsa.PublicKey(msg['available_cli']['n'], 
                    msg['available_cli']['e'])
            clis.append(k)

    elif msg['type'] == 'message':
        # check if it's only a partial message
        try:
            if 'from' in msg:
                if msg['id'] not in messages:
                    messages[msg['id']] = {}
                messages[msg['id']][msg['from']] = msg['contents']
                contents = None
                for i in range(0, 3):
                    if i in messages[msg['id']]:
                        if contents is None:
                            contents = messages[msg['id']][i].encode('latin')
                        else:
                            contents += messages[msg['id']][i].encode('latin')
                if len(contents) == 256:
                    m = rsa.decrypt(contents, priv).decode()
                else:
                    m = '                                      '
            else:
                contents = msg['contents'].encode('latin')
                m = rsa.decrypt(contents, priv).decode()
            if m is not None:
                print(f"\033[F\033[1G<< {m}\n> ")
        except Exception as e:
            print(f'Failed to receive message due to {e}')


def on_error(ws, error):
    print(error)


def on_close(ws, close_status_code, close_msg):
    print(f"### closed. reason: {close_msg} ###")


def on_open(ws):
    def run(*args):
        init_send(ws)
        running = True
        while running:
            outbound = input(">")
            if outbound == 'q':
                running = False
            else:
                send(ws, outbound)
        ws.close()
    _thread.start_new_thread(run, ())


def open_socket(url, port):
    ws = websocket.WebSocketApp(f"ws://{url}:{port}",
        on_open=on_open,
        on_message=on_message,
        on_error=on_error,
        on_close=on_close)
    ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})


if __name__ == "__main__":
    with open('src/config.yml', 'r') as f:
        conf = yaml.safe_load(f)

    e = input("Are you outside of the network? (Y/n)")
    if e == 'y' or e == 'Y' or e == 'yes' or e == 'Yes':
        ext = True
        ws_conf = conf['ws']['ext']
    else:
        ws_conf = conf['ws']['int']
    print('Generating keys...')
    pub, priv = rsa.newkeys(2048)
    print('Keys generated')
    open_socket(ws_conf['url'], ws_conf['port'])