from __future__ import annotations
from dataclasses import dataclass
from functools import reduce
from hashlib import sha256
from time import time
from nacl.public import Box, SealedBox
from nacl.signing import SigningKey, VerifyKey, SignedMessage
from queue import SimpleQueue
from random import randint
from secrets import token_bytes
from typing import Protocol
import readline
import struct


"""Proof-of-concept gossip protocol implementation. To use in a real
    application, import the module; create new handlers as necessary for
    the duck typed protocols SupportsSendAndDeliverMessage,
    SupportsHandleMessage, and SupportsHandleAction; and register the
    handlers as shown in main. Note that Node.from_seed is used only for
    the active node(s) while Node.__init__ is can be used for neighbors
    for which we know only the public key/address. Most models include
    a data or metadata property for extensibility. The run_tick,
    format_address, and action_count functions can be used without
    modification. Debug message handling can be swapped out from print
    to a custom function using deregister_debug_handler(print) and
    register_debug_handler(custom_func), e.g. to write to a log file.
"""


def license() -> str:
    """Copyleft (c) 2022 k98kurz

        Permission to use, copy, modify, and/or distribute this software
        for any purpose with or without fee is hereby granted, provided
        that the above copyleft notice and this permission notice appear in
        all copies.

        THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
        WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
        WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
        AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR
        CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
        OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
        NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
        CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
    """
    return license.__doc__


# global config and toggle/utility functions
ENABLE_DEBUG = True
DISPLAY_SHORT_ADDRESSES = True
SIGN_MESSAGES = True
ENCRYPT_MESSAGES = True
MESSAGE_TTL = 300
DEBUG_HANDLERS = [print]

def format_address(address: bytes) -> str:
    global DISPLAY_SHORT_ADDRESSES
    return address.hex()[:8] if DISPLAY_SHORT_ADDRESSES else address.hex()

def toggle_short_address() -> bool:
    global DISPLAY_SHORT_ADDRESSES
    DISPLAY_SHORT_ADDRESSES = not DISPLAY_SHORT_ADDRESSES
    return DISPLAY_SHORT_ADDRESSES

def debug(msg: str):
    """Pass debug messages to all debug message handlers."""
    global ENABLE_DEBUG, DEBUG_HANDLERS
    if ENABLE_DEBUG:
        for d in DEBUG_HANDLERS:
            d(msg)

def register_debug_handler(c: function) -> None:
    """Register a new function for handling debug messages."""
    if not callable(c):
        raise TypeError('Can only register callables as debug handlers.')
    global DEBUG_HANDLERS
    if c not in DEBUG_HANDLERS:
        DEBUG_HANDLERS.append(c)

def unregister_debug_handler(c: function) -> None:
    """Unregister a function from handling debug messages."""
    if not callable(c):
        raise TypeError('Can only deregister callables as debug handlers.')
    global DEBUG_HANDLERS
    if c in DEBUG_HANDLERS:
        DEBUG_HANDLERS.remove(c)

def toggle_debug() -> bool:
    global ENABLE_DEBUG
    ENABLE_DEBUG = not ENABLE_DEBUG
    return ENABLE_DEBUG

def toggle_sign_all_messages() -> bool:
    global SIGN_MESSAGES
    SIGN_MESSAGES = not SIGN_MESSAGES
    return SIGN_MESSAGES

def toggle_encrypt_messages() -> bool:
    global ENCRYPT_MESSAGES
    ENCRYPT_MESSAGES = not ENCRYPT_MESSAGES
    return ENCRYPT_MESSAGES


@dataclass
class Connection:
    """Connection model represent an edge connecting two Nodes together."""
    nodes: set[Node]
    data: dict

    def __init__(self, nodes: list[Node]) -> None:
        if type(nodes) is not list or len(nodes) != 2:
            raise Exception('a Connection must connect exactly 2 nodes')
        self.nodes = set(nodes)
        self.data = {}

    def __hash__(self) -> int:
        """Enable inclusion in sets."""
        node_list = list(self.nodes)
        node_list.sort()
        return hash(node_list[0].address + node_list[1].address)


@dataclass
class Message:
    """Message model contains the source, destination, content, and
        optional signature and metadata.
    """
    src: bytes
    dst: bytes
    ts: int
    msg: bytes
    sig: bytes
    metadata: dict

    def __init__(self, src: bytes, dst: bytes, msg: bytes, ts: int = None, sig: bytes = None) -> None:
        self.src = src
        self.dst = dst
        self.ts = int(time()) if ts is None else ts
        self.msg = msg
        self.sig = sig
        self.metadata = {}

    def __repr__(self) -> str:
        return f"{format_address(self.src)}->{format_address(self.dst)}: {format_address(sha256(self.msg).digest())}"

    def __bytes__(self) -> bytes:
        return self.src + self.dst + self.msg

    def __hash__(self) -> int:
        """Enable inclusion in sets."""
        return hash(bytes(self))

    def pack(self) -> bytes:
        """Pack the data with struct."""
        if self.sig is not None:
            fstr = '!32s32si64s' + str(len(self.msg)) + 's'
            return struct.pack(fstr, self.dst, self.src, self.ts, self.sig, self.msg)
        else:
            fstr = '!32s32si' + str(len(self.msg)) + 's'
            return struct.pack(fstr, self.dst, self.src, self.ts, self.msg)

    @classmethod
    def unpack(cls, packed: bytes) -> Message:
        """Unpack the data with struct."""
        global SIGN_MESSAGES
        if SIGN_MESSAGES:
            fstr = '!32s32si64s' + str(len(packed) - 128) + 's'
            (dst, src, ts, sig, msg) = struct.unpack(fstr, packed)
            return Message(src, dst, msg, ts, sig)
        else:
            fstr = '!32s32si' + str(len(packed) - 64) + 's'
            (dst, src, ts, msg) = struct.unpack(fstr, packed)
            return Message(src, dst, msg, ts)

    def sign(self, skey: SigningKey) -> SignedMessage:
        """Generate a signature for the message."""
        sig = skey.sign(bytes(self))
        self.sig = sig[:64]
        return sig

    def verify(self) -> bool:
        """Verify the message signature"""
        try:
            vkey = VerifyKey(self.src)
            sig = SignedMessage(self.sig + bytes(self))
            vkey.verify(sig)
            return True
        except:
            return False

    def encrypt(self, skey: SigningKey) -> None:
        """Encrypt the message by the sender."""
        if bytes(skey.verify_key) != self.src:
            raise ValueError('Must use the skey of the sender to encrypt.')

        privk = skey.to_curve25519_private_key()
        pubk = VerifyKey(self.dst).to_curve25519_public_key()
        box = Box(privk, pubk)
        self.msg = bytes(box.encrypt(self.msg))

    def decrypt(self, skey: SigningKey) -> None:
        """Decrypt the message by the receiver."""
        if bytes(skey.verify_key) != self.dst:
            raise ValueError('Must use the skey of the receiver to decrypt.')

        privk = skey.to_curve25519_private_key()
        pubk = VerifyKey(self.src).to_curve25519_public_key()
        box = Box(privk, pubk)
        self.msg = box.decrypt(self.msg)

    def seal(self) -> None:
        """Encrypt using ephemeral ECDHE."""
        sealed_box = SealedBox(VerifyKey(self.dst).to_curve25519_public_key())
        self.msg = sealed_box.encrypt(self.msg)

    def unseal(self, skey: SigningKey) -> None:
        """Decrypt using ephemeral ECDHE."""
        if bytes(skey.verify_key) != self.dst:
            raise ValueError('Must use the skey of the receiver to decrypt.')

        privk = skey.to_curve25519_private_key()
        sealed_box = SealedBox(privk)
        self.msg = sealed_box.decrypt(self.msg)


@dataclass
class Action:
    """Action model contains the name and data for an action a Node will
        take by passing to the registered action handler.
    """
    name: str
    data: dict

    def __init__(self, name: str, data: dict) -> None:
        self.name = name
        self.data = data


class SupportsSendAndDeliverMessage(Protocol):
    """Duck type protocol for message sender."""
    def send(self, msg: Message) -> None:
        ...

    def deliver(self) -> None:
        ...


class SupportsHandleMessage(Protocol):
    """Duck type protocol for incoming message handler."""
    def handle(self, msg: Message) -> None:
        ...


class SupportsHandleAction(Protocol):
    """Duck type protocol for action handler."""
    def handle(self, action: dict) -> None:
        ...


@dataclass
class Node:
    """The core model representing a Node and handling its Connections,
        Actions, and Messages. Invoke with Node(address) for neighbors
        and Node.from_seed(seed) for an active node. Optional data
        property for extensibility. The address is the public key bytes
        of the node for when SIGN_MESSAGES is set.
    """
    address: bytes
    msgs_seen: set[bytes]
    connections: set[Connection]
    data: dict
    _seed: bytes
    _skey: SigningKey
    _vkey: VerifyKey
    _inbound: SimpleQueue
    _outbound: SimpleQueue
    _actions: SimpleQueue
    _message_sender: SupportsSendAndDeliverMessage
    _message_handler: SupportsHandleMessage
    _action_handler: SupportsHandleAction

    def __init__(self, address: bytes) -> None:
        """Create a node from its address (public key bytes)."""
        self.address = address
        self.msgs_seen = set()
        self.connections = set()
        self.data = {}
        self._vkey = VerifyKey(address)
        self._seed = None
        self._skey = None
        self._inbound = SimpleQueue()
        self._outbound = SimpleQueue()
        self._actions = SimpleQueue()

    @classmethod
    def from_seed(cls, seed: bytes):
        """Create a node from a seed filling out _skey."""
        skey = SigningKey(seed)
        node = cls(bytes(skey.verify_key))
        node._skey = skey
        node._seed = seed
        return node

    def __hash__(self) -> int:
        """Enable inclusion in sets."""
        return hash(self.address)

    def __lt__(self, other: Node) -> bool:
        return self.address < other.address

    def __repr__(self) -> str:
        if self._seed is not None:
            return "{'address': '" + format_address(self.address) + "','seed':'" + self._seed.hex() + "}"
        else:
            return "{'address': '" + format_address(self.address) + "'}"

    def register_message_sender(self, sndr: SupportsSendAndDeliverMessage) -> None:
        """Register the message sender."""
        if not hasattr(sndr, 'send') or not callable(sndr.send):
            raise TypeError('sndr must fulfill SupportsSendAndDeliverMessage duck type')
        self._message_sender = sndr

    def register_message_handler(self, hndlr: SupportsHandleMessage) -> None:
        """Register the incoming message handler."""
        if not hasattr(hndlr, 'handle') or not callable(hndlr.handle):
            raise TypeError('hndlr must fulfill SupportsHandleMessage duck type')
        self._message_handler = hndlr

    def register_action_handler(self, hndlr: SupportsHandleAction) -> None:
        """Register the action handler."""
        if not hasattr(hndlr, 'handle') or not callable(hndlr.handle):
            raise TypeError('hndlr must fulfill SupportsHandleAction duck type')
        self._action_handler = hndlr

    def add_connection(self, connection: Connection) -> None:
        """Add the specified connection."""
        if not isinstance(connection, Connection):
            raise TypeError('connection must be a Connection')
        self.connections.add(connection)

    def drop_connection(self, connection: Connection) -> None:
        """Drop the specified connection."""
        if not isinstance(connection, Connection):
            raise TypeError('connection must be a Connection')
        self.connections.remove(connection)

    def count_connections(self) -> int:
        return len(self.connections)

    def receive_message(self, message: Message):
        """Queue up an incoming message if its signature is valid or
            ignored.
        """
        if not isinstance(message, Message):
            raise TypeError('message must be a Message')

        global SIGN_MESSAGES, ENCRYPT_MESSAGES, MESSAGE_TTL

        if int(time()) > (message.ts + MESSAGE_TTL):
            debug("Node.receive_message: old message discarded")
        elif message.sig is not None:
            if message.verify():
                if ENCRYPT_MESSAGES:
                    message.unseal(self._skey)
                self._inbound.put(message)
            else:
                debug("Node.receive_message: message signature failed verification")
        elif SIGN_MESSAGES:
            debug("Node.receive_message: unsigned message rejected")
        else:
            if ENCRYPT_MESSAGES:
                message.unseal(self._skey)
            self._inbound.put(message)

    def send_message(self, dst: bytes, msg: bytes):
        """Queue up an outgoing message. Sign if necessary and possible."""
        if type(dst) is not bytes:
            raise TypeError("dst must be bytes")
        if type(msg) is not bytes:
            raise TypeError("msg must be bytes")

        message = Message(self.address, dst, msg)

        global ENCRYPT_MESSAGES
        if ENCRYPT_MESSAGES:
            message.seal()

        if self._skey is not None:
            message.sign(self._skey)

        if len(self.connections):
            if len([c for c in self.connections if dst in [n.address for n in c.nodes]]):
                self._outbound.put(message)
            else:
                debug("cannot deliver message due to lack of connection")
        else:
            self._outbound.put(message)

    def queue_action(self, act: Action) -> None:
        """Queue an action to be processed by the action handler."""
        if not isinstance(act, Action):
            raise TypeError('act must be an Action')
        self._actions.put(act)

    def process(self):
        """Process actions for this node once."""
        if self._outbound.qsize() > 0 and self._message_sender is not None:
            self._message_sender.send(self._outbound.get())
        if self._inbound.qsize() > 0 and self._message_handler is not None:
            self._message_handler.handle(self._inbound.get())
        if self._actions.qsize() > 0 and self._action_handler is not None:
            self._action_handler.handle(self._actions.get())

    def action_count(self):
        """Count the size of pending messages and actions."""
        return self._outbound.qsize() + self._inbound.qsize() + self._actions.qsize()


@dataclass
class MessageSender:
    """Example message sender that does not use a network stack."""
    nodes: set[Node]
    message_queue: SimpleQueue
    dead_letters: list
    data: dict

    def __init__(self) -> None:
        self.nodes = set()
        self.message_queue = SimpleQueue()
        self.dead_letters = []
        self.data = {}

    def register_node(self, node: Node) -> None:
        """Register the specified node for message delivery."""
        if not isinstance(node, Node):
            raise TypeError('node must be a Node')
        self.nodes.add(node)
        self.check_dead_letters(node)

    def register_nodes(self, nodes: list[Node]) -> None:
        """Register the specified list of nodes for message delivery."""
        if not isinstance(nodes, list):
            raise TypeError('nodes must be a list of Nodes')
        for n in nodes:
            self.register_node(n)

    def send(self, msg: Message) -> None:
        """Queue up the specified message for delivery."""
        if not isinstance(msg, Message):
            raise TypeError('msg must be a Message')
        debug(f"MessageSender.send(): {msg}")
        self.message_queue.put(msg)

    def check_queue(self) -> int:
        """Return pending message queue size."""
        return self.message_queue.qsize()

    def deliver(self) -> None:
        """Deliver all pending messages that can be delivered."""
        while self.check_queue() > 0:
            msg = self.message_queue.get()
            found = False

            for n in self.nodes:
                if n.address == msg.dst:
                    n.receive_message(msg)
                    found = True
                    debug(f'MessageSender.deliver(): delivered {format_address(sha256(msg.msg).digest())} to {format_address(n.address)}')

            if not found:
                self.dead_letters.append(msg)
                debug('MessageSender.deliver(): dead letter')

    def check_dead_letters(self, node: Node) -> None:
        """Goes through the dead letters to try to deliver to a newly
            registered node.
        """
        if not isinstance(node, Node):
            raise TypeError('node must be a Node')

        delivered = []
        for i, l in enumerate(self.dead_letters):
            if l.dst == node.address:
                node.receive_message(l)
                delivered.append(l)
        self.dead_letters[:] = [l for l in self.dead_letters if l not in delivered]


@dataclass
class MessageHandler:
    """Example message handler. Optional data property for extensibility."""
    nodes: set[Node]
    data: dict

    def __init__(self) -> None:
        self.nodes = set()
        self.data = {}

    def register_node(self, node: Node) -> None:
        """Register the specified node for handling incoming messages."""
        if not isinstance(node, Node):
            raise TypeError('node must be a Node')
        self.nodes.add(node)

    def register_nodes(self, nodes: list[Node]) -> None:
        """Register the specified list of nodes for handling incoming
            messages.
        """
        if not isinstance(nodes, list):
            raise TypeError('nodes must be a list of Nodes')
        for n in nodes:
            self.register_node(n)

    def handle(self, msg: Message) -> None:
        """Handle an incoming message."""
        if not isinstance(msg, Message):
            raise TypeError('msg must be a Message')

        debug(f'MessageHandler.handle(): {msg}')

        if msg.dst in [n.address for n in self.nodes]:
            n = [n for n in self.nodes if n.address == msg.dst][0]
            n.queue_action(Action('store_and_forward', {"msg": msg.msg}))
        else:
            debug('MessageHandler.handle(): message dropped')


@dataclass
class ActionHandler:
    """Example action handler."""
    node: Node
    other_nodes: set[Node]

    def __init__(self, node: Node, other_nodes: list[Node]) -> None:
        self.node = node
        self.other_nodes = set(other_nodes)

    def handle(self, act: Action) -> None:
        """Handle an action. Limited to store_and_forward action."""
        if act.name == 'store_and_forward':
            if sha256(act.data['msg']).digest() not in self.node.msgs_seen:
                # store
                self.node.msgs_seen.add(sha256(act.data['msg']).digest())
                debug(f"ActionHandler.handle(): store_and_forward [{act.data['msg'].hex()}]")

                # forward
                if self.node.count_connections() > 0:
                    # forward to all connected nodes
                    for c in self.node.connections:
                        n = [n for n in c.nodes if n is not self.node][0]
                        self.node.send_message(n.address, act.data['msg'])
                else:
                    # forward to 2 nodes at random
                    n1 = list(self.other_nodes)[randint(0, len(self.other_nodes)-1)]
                    n2 = list(self.other_nodes)[randint(0, len(self.other_nodes)-1)]
                    self.node.send_message(n1.address, act.data['msg'])
                    self.node.send_message(n2.address, act.data['msg'])
            else:
                debug(f"ActionHandler.handle(): store_and_forward skipped for seen message")


def run_tick(nodes: list[Node], msg_sender: MessageSender):
    """Run the process for all nodes, then deliver all pending messages."""
    for n in nodes:
        n.process()
    msg_sender.deliver()


def action_count(nodes: list[Node]):
    """Returns a count of all pending actions and messages."""
    if not isinstance(nodes, list):
        raise TypeError('nodes must be list of Nodes')
    return reduce(lambda c, n: c + n.action_count(), nodes, 0)


def main():
    # create some nodes
    nodes = [Node.from_seed(token_bytes(32)) for i in range(16)]

    # create handlers
    msg_handler = MessageHandler()
    msg_handler.register_nodes(nodes)
    msg_sender = MessageSender()
    msg_sender.register_nodes(nodes)

    # register handlers
    for n in nodes:
        n.register_message_handler(msg_handler)
        n.register_message_sender(msg_sender)
        n.register_action_handler(ActionHandler(n, [on for on in nodes if on is not n]))

    # flag for exit
    end_signal = False

    # access global variables
    global SIGN_MESSAGES, ENCRYPT_MESSAGES

    while not end_signal:
        data = input("$: ")
        command = data.split(' ')[0].strip()
        data = ' '.join(data.split(' ')[1:]).strip()

        if command in ('quit', 'q'):
            end_signal = True
        elif command in ('list', 'nodes', 'l', 'n', 'ln'):
            for n in nodes:
                print(f"{format_address(n.address)}: {[format_address(m) for m in n.msgs_seen]}")
        elif command in ('listcon', 'connections', 'lc'):
            connections = set()
            for n in nodes:
                connections = connections.union(n.connections)
            for c in connections:
                cnodes = list(c.nodes)
                print(f"{format_address(cnodes[0].address)} - {format_address(cnodes[1].address)}")
        elif command in ('c', 'connect'):
            for n in nodes:
                others = [o for o in nodes if o is not n]
                for i in range(3):
                    o = others[randint(0, len(others)-1)]
                    n.add_connection(Connection([n, o]))
                    o.add_connection(Connection([n, o]))
        elif command in ('message', 'm'):
            src = nodes[randint(0, len(nodes)-1)]
            message = Message(src.address, src.address, bytes(data, 'utf-8'))
            if ENCRYPT_MESSAGES:
                message.seal()
            if SIGN_MESSAGES:
                message.sign(src._skey)
            src.receive_message(message)
        elif command in ('d', 'debug'):
            print("debug enabled" if toggle_debug() else "debug disabled")
        elif command in ('s', 'short'):
            print("short addresses enabled" if toggle_short_address() else "short addresses disabled")
        elif command in ('r', 'run'):
            while action_count(nodes) > 0:
                run_tick(nodes, msg_sender)
        elif command in ('h', 'help', '?'):
            print("options:\t[l|ln|nodes|list] to list nodes and messages seen by each")
            print("\t\t[m|message] {str} to send a message")
            print("\t\t[c|connect] to connect nodes together randomly")
            print("\t\t[lc|listcon|connections] list all connections")
            print("\t\t[q|quit] to end")
            print("\t\t[h|help|?] display this text")
            print("\t\t[d|debug] to toggle debug messages")
            print("\t\t[s|short] to toggle displaying short address format")
            print("\t\t[r|run] to run until no pending actions remain")
            print("\t\tanything else to process a tick")
        else:
            run_tick(nodes, msg_sender)


if __name__ == '__main__':
    main()