#!/usr/bin/env python3

import random # For simulation, not cryptography!
import math
import sys
import os
import logging
import resource
import sympy

sys.path.append("..")

import network
import dirauth
import relay
import client

class BandwidthMeasurer:
    def __init__(self, numdirauths, numrelays, numclients):
        # Start some dirauths
        self.dirauthaddrs = []
        self.dirauths = []
        for i in range(numdirauths):
            dira = dirauth.DirAuth(i, numdirauths)
            self.dirauths.append(dira)
            self.dirauthaddrs.append(dira.netaddr)

        # Start some relays
        self.relays = []
        for i in range(numrelays):
            self.startrelay()

        # The fallback relays are a hardcoded list of a small fraction
        # of the relays, used by clients for bootstrapping
        numfallbackrelays = 1
        fallbackrelays = self.relays[0:1]
        for r in fallbackrelays:
            r.set_is_fallbackrelay()
        network.thenetwork.setfallbackrelays(fallbackrelays)

        # Tick the epoch to build the first consensus
        network.thenetwork.nextepoch()

        # Start some clients
        self.clients = []
        for i in range(numclients):
            self.startclient()

        # Throw away all the performance statistics to this point
        for d in self.dirauths: d.perfstats.reset()
        for r in self.relays: r.perfstats.reset()
        # The clients' stats are already at 0, but they have the
        # "bootstrapping" flag set, which we want to keep, so we
        # won't reset them.

        self.allcircs = []

        # Tick the epoch to bootstrap the clients
        network.thenetwork.nextepoch()

    def startrelay(self):
        # Relay bandwidths (at least the ones fast enough to get used)
        # in the live Tor network (as of Dec 2019) are well approximated
        # by (200000-(200000-25000)/3*log10(x)) where x is a
        # uniform integer in [1,2500]
        x = random.randint(1,2500)
        bw = int(200000-(200000-25000)/3*math.log10(x))
        self.relays.append(relay.Relay(self.dirauthaddrs, bw, 0))

    def stoprelay(self):
        self.relays[1].terminate()
        del self.relays[1]

    def startclient(self):
        self.clients.append(client.Client(self.dirauthaddrs))

    def stopclient(self):
        self.clients[0].terminate()
        del self.clients[0]

    def buildcircuit(self):
        bwm.allcircs.append(bwm.clients[0].channelmgr.new_circuit())

    def getstats(self):

        # gather stats
        totsent = 0
        totrecv = 0
        totbytes = 0
        dirasent = 0
        dirarecv = 0
        dirabytes = 0
        relaysent = 0
        relayrecv = 0
        relaybytes = 0
        clisent = 0
        clirecv = 0
        clibytes = 0
        for d in self.dirauths:
            logging.debug("%s", d.perfstats)
            dirasent += d.perfstats.bytes_sent
            dirarecv += d.perfstats.bytes_received
            dirabytes += d.perfstats.bytes_sent + d.perfstats.bytes_received
        totsent += dirasent
        totrecv += dirarecv
        totbytes += dirabytes
        for r in self.relays:
            logging.debug("%s", r.perfstats)
            relaysent += r.perfstats.bytes_sent
            relayrecv += r.perfstats.bytes_received
            relaybytes += r.perfstats.bytes_sent + r.perfstats.bytes_received
        totsent += relaysent
        totrecv += relayrecv
        totbytes += relaybytes
        for c in self.clients:
            logging.debug("%s", c.perfstats)
            clisent += c.perfstats.bytes_sent
            clirecv += c.perfstats.bytes_received
            clibytes += c.perfstats.bytes_sent + c.perfstats.bytes_received
        totsent += clisent
        totrecv += clirecv
        totbytes += clibytes
        logging.info("DirAuths sent=%s recv=%s bytes=%s" % \
                (dirasent, dirarecv, dirabytes))
        logging.info("Relays sent=%s recv=%s bytes=%s" % \
                (relaysent, relayrecv, relaybytes))
        logging.info("Client sent=%s recv=%s bytes=%s" % \
                (clisent, clirecv, clibytes))
        logging.info("Total sent=%s recv=%s bytes=%s" % \
                (totsent, totrecv, totbytes))

        # Reset bootstrap flag
        for d in self.dirauths: d.perfstats.is_bootstrapping = False
        for r in self.relays: r.perfstats.is_bootstrapping = False
        for c in self.clients: c.perfstats.is_bootstrapping = False

        return (dirabytes, relaybytes, clibytes)

    def endepoch(self):

        # Close circuits
        for c in self.allcircs:
            c.close()
        self.allcircs = []

        # Reset stats
        for d in self.dirauths: d.perfstats.reset()
        for r in self.relays: r.perfstats.reset()
        for c in self.clients: c.perfstats.reset()

        network.thenetwork.nextepoch()


if __name__ == '__main__':
    # Args: womode snipauthmode numrelays randseed
    if len(sys.argv) != 5:
        sys.stderr.write("Usage: womode snipauthmode numrelays randseed\n")
        sys.exit(1)

    womode = network.WOMode[sys.argv[1].upper()]
    snipauthmode = network.SNIPAuthMode[sys.argv[2].upper()]
    numrelays = int(sys.argv[3])
    randseed = int(sys.argv[4])

    # Use symbolic byte counter mode
    network.symbolic_byte_counters = True

    # Seed the PRNG.  On Ubuntu 18.04, this in fact makes future calls
    # to (non-cryptographic) random numbers deterministic.  On Ubuntu
    # 16.04, it does not.
    random.seed(randseed)

    loglevel = logging.INFO
    # Uncomment to see all the debug messages
    # loglevel = logging.DEBUG

    logging.basicConfig(level=loglevel,
            format="%(asctime)s:%(levelname)s:%(message)s")

    logging.info("Starting simulation")

    # Set the Walking Onions style to use
    network.thenetwork.set_wo_style(womode, snipauthmode)

    bwm = BandwidthMeasurer(9, numrelays, 0)
    stats = dict()

    logging.info("R_N = %d, R_B = 0, C_N = 0, C_B = 0, circs = 0", numrelays)
    stats[(numrelays, 0, 0, 0, 0)] = bwm.getstats()
    # Bootstrap one relay
    bwm.startrelay()
    bwm.endepoch()
    logging.info("R_N = %d, R_B = 1, C_N = 0, C_B = 0, circs = 0", numrelays)
    stats[(numrelays, 1, 0, 0, 0)] = bwm.getstats()
    # Bootstrap one client
    bwm.stoprelay()
    bwm.startclient()
    bwm.endepoch()
    logging.info("R_N = %d, R_B = 0, C_N = 0, C_B = 1, circs = 0", numrelays)
    stats[(numrelays, 0, 0, 1, 0)] = bwm.getstats()
    # No changes, so the client is now not bootstrapping
    bwm.endepoch()
    logging.info("R_N = %d, R_B = 0, C_N = 1, C_B = 0, circs = 0", numrelays)
    stats[(numrelays, 0, 1, 0, 0)] = bwm.getstats()
    # No more bootstrapping, but build one circuit
    bwm.buildcircuit()
    logging.info("R_N = %d, R_B = 0, C_N = 1, C_B = 0, circs = 1", numrelays)
    stats[(numrelays, 0, 1, 0, 1)] = bwm.getstats()
    bwm.endepoch()
    # No more bootstrapping, but build two circuits
    bwm.buildcircuit()
    bwm.buildcircuit()
    logging.info("R_N = %d, R_B = 0, C_N = 1, C_B = 0, circs = 2", numrelays)
    stats[(numrelays, 0, 1, 0, 2)] = bwm.getstats()
    bwm.endepoch()

    print("\n")
    print('Total relay bytes:')
    print('  R_N * (', stats[(numrelays, 0, 0, 0, 0)][1]/numrelays, ')')
    print('+ R_B * (', stats[(numrelays, 1, 0, 0, 0)][1] - stats[(numrelays, 0, 0, 0, 0)][1], ')')
    print('+ C_N * (', stats[(numrelays, 0, 1, 0, 0)][1] - stats[(numrelays, 0, 0, 0, 0)][1], ')')
    print('+ C_B * (', stats[(numrelays, 0, 0, 1, 0)][1] - stats[(numrelays, 0, 0, 0, 0)][1], ')')
    print('+ circ * (', stats[(numrelays, 0, 1, 0, 1)][1] - stats[(numrelays, 0, 1, 0, 0)][1], ')')
    print('  check   ', stats[(numrelays, 0, 1, 0, 2)][1] - stats[(numrelays, 0, 1, 0, 1)][1])

    print("\n")
    print('Total client bytes:')
    print('  R_N * (', stats[(numrelays, 0, 0, 0, 0)][2]/numrelays, ')')
    print('+ R_B * (', stats[(numrelays, 1, 0, 0, 0)][2] - stats[(numrelays, 0, 0, 0, 0)][2], ')')
    print('+ C_N * (', stats[(numrelays, 0, 1, 0, 0)][2] - stats[(numrelays, 0, 0, 0, 0)][2], ')')
    print('+ C_B * (', stats[(numrelays, 0, 0, 1, 0)][2] - stats[(numrelays, 0, 0, 0, 0)][2], ')')
    print('+ circ * (', stats[(numrelays, 0, 1, 0, 1)][2] - stats[(numrelays, 0, 1, 0, 0)][2], ')')
    print('  check   ', stats[(numrelays, 0, 1, 0, 2)][2] - stats[(numrelays, 0, 1, 0, 1)][2])