import pickle
import math

import network

# Set this to True if you want the bytes sent and received to be added
# symbolically, in terms of the numbers of each type of network message.
# You will need sympy installed for this to work.
symbolic_byte_counters = False

try:
    import sympy
except:
    pass

class NetMsg:
    """The parent class of network messages.  Subclass this class to
    implement specific kinds of network messages."""

    def size(self):
        """Return the size of this network message.  For now, just
        pickle it and return the length of that.  There's some
        unnecessary overhead in this method; if you want specific
        messages to have more accurate sizes, override this method in
        the subclass.  Alternately, if symbolic_byte_counters is set,
        return a symbolic representation of the message size instead, so
        that the total byte counts will clearly show how many of each
        message type were sent and received."""
        if symbolic_byte_counters:
            sz = sympy.symbols(type(self).__name__)
        else:
            sz = len(pickle.dumps(self))
            # logging.info("%s size %d", type(self).__name__, sz)

        return sz


class StringNetMsg(NetMsg):
    """Send an arbitratry string as a NetMsg."""
    def __init__(self, data):
        self.data = data

    def __str__(self):
        return self.data.__str__()

#########################
# DIRAUTH
#########################

class DirAuthNetMsg(NetMsg):
    """The subclass of NetMsg for messages to and from directory
    authorities."""


class DirAuthUploadDescMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for uploading a relay
    descriptor."""

    def __init__(self, desc):
        self.desc = desc


class DirAuthDelDescMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for deleting a relay
    descriptor."""

    def __init__(self, desc):
        self.desc = desc


class DirAuthGetConsensusMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for fetching the consensus."""


class DirAuthConsensusMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for returning the consensus."""

    def __init__(self, consensus):
        self.consensus = consensus


class DirAuthGetConsensusDiffMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for fetching the consensus, if the
    requestor already has the previous consensus."""


class DirAuthConsensusDiffMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for returning the consensus, if the
    requestor already has the previous consensus.  We don't _actually_
    produce the diff at this time; we just charge fewer bytes for this
    message."""

    def __init__(self, consensus):
        self.consensus = consensus

    def size(self):
        if symbolic_byte_counters:
            return super().size()
        return math.ceil(DirAuthConsensusMsg(self.consensus).size() \
                            * network.P_Delta)

class DirAuthGetENDIVEMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for fetching the ENDIVE."""


class DirAuthENDIVEMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for returning the ENDIVE."""

    def __init__(self, endive):
        self.endive = endive


class DirAuthGetENDIVEDiffMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for fetching the ENDIVE, if the
    requestor already has the previous ENDIVE."""


class DirAuthENDIVEDiffMsg(DirAuthNetMsg):
    """The subclass of DirAuthNetMsg for returning the ENDIVE, if the
    requestor already has the previous consensus.  We don't _actually_
    produce the diff at this time; we just charge fewer bytes for this
    message in Merkle mode.  In threshold signature mode, we would still
    need to download at least the new signatures for every SNIP in the
    ENDIVE, so for now, just assume there's no gain from ENDIVE diffs in
    threshold signature mode."""

    def __init__(self, endive):
        self.endive = endive

    def size(self):
        if symbolic_byte_counters:
            return super().size()
        if network.thenetwork.snipauthmode == \
                network.SNIPAuthMode.THRESHSIG:
            return DirAuthENDIVEMsg(self.endive).size()
        return math.ceil(DirAuthENDIVEMsg(self.endive).size() \
                            * network.P_Delta)

#########################
# RELAY
#########################

class RelayNetMsg(NetMsg):
    """The subclass of NetMsg for messages between relays and either
    relays or clients."""


class RelayGetConsensusMsg(RelayNetMsg):
    """The subclass of RelayNetMsg for fetching the consensus. Sent by
    clients to relays."""


class RelayConsensusMsg(RelayNetMsg):
    """The subclass of RelayNetMsg for returning the consensus from
    relays to clients, in response to RelayGetConsensusMsg."""

    def __init__(self, consensus):
        self.consensus = consensus


class RelayGetDescMsg(RelayNetMsg):
    """The subclass of RelayNetMsg sent by clients to their guards for
    retrieving the guard's current descriptor."""


class RelayDescMsg(RelayNetMsg):
    """The subclass of RelayNetMsg sent by guards to clients for
    reporting their current descriptor."""

    def __init__(self, desc):
        self.desc = desc


class RelayGetConsensusDiffMsg(RelayNetMsg):
    """The subclass of RelayNetMsg for fetching the consensus, if the
    requestor already has the previous consensus. Sent by clients to
    relays."""


class RelayConsensusDiffMsg(RelayNetMsg):
    """The subclass of RelayNetMsg for returning the consensus, if the
    requestor already has the previous consensus.  We don't _actually_
    produce the diff at this time; we just charge fewer bytes for this
    message. Sent by relays to clients in response to
    RelayGetConsensusDiffMsg."""

    def __init__(self, consensus):
        self.consensus = consensus

    def size(self):
        if symbolic_byte_counters:
            return super().size()
        return math.ceil(RelayConsensusMsg(self.consensus).size() \
                            * network.P_Delta)


class RelayRandomHopMsg(RelayNetMsg):
    """A message used for testing, that hops from relay to relay
    randomly until its TTL expires."""

    def __init__(self, ttl):
        self.ttl = ttl

    def __str__(self):
        return "RandomHop TTL=%d" % self.ttl


class CircuitCellMsg(RelayNetMsg):
    """Send a message tagged with a circuit id. This is the container
    class for all RelayCell messages."""

    def __init__(self, circuitid, cell):
        self.circid = circuitid
        self.cell = cell

    def __str__(self):
        return "C%d:%s" % (self.circid, self.cell)

    def size(self):
        # circuitids are 4 bytes
        return 4 + self.cell.size()


# It is intentional that VanillaCreateCircuitMsg is a RelayNetMsg and
# not a RelayCell.  This is the message that _creates_ the circuit, so
# it can't be sent as a cell _within_ the circuit.
class VanillaCreateCircuitMsg(RelayNetMsg):
    """The message for requesting circuit creation in Vanilla Onion
    Routing."""

    def __init__(self, circid, ntor_request):
        self.circid = circid
        self.ntor_request = ntor_request


# It is intentional that TelescopingCreateCircuitMsg is a RelayNetMsg and
# not a RelayCell.  This is the message that _creates_ the circuit, so
# it can't be sent as a cell _within_ the circuit.
class TelescopingCreateCircuitMsg(RelayNetMsg):
    """The message for requesting circuit creation in Telescoping Onion
    Routing."""

    def __init__(self, circid, ntor_request):
        self.circid = circid
        self.ntor_request = ntor_request


class SinglePassCreateCircuitMsg(RelayNetMsg):
    """The message for requesting circuit creation in Single Pass Onion
    Routing. This is used to extend a Single-Pass circuit by clients."""

    def __init__(self, circid, ntor_request, client_path_selection_key,
            ttl=2):
        self.circid = circid
        self.ntor_request = ntor_request
        self.clipathselkey = bytes(client_path_selection_key)
        self.ttl = ttl