#!/usr/bin/env python3

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

sys.path.append("..")

import network
import dirauth
import relay
import client

from bwparser import JansenBandwidthParser, KomloBandwidthParser

class BandwidthMeasurer:
    def __init__(self, total_numrelays, bw_parser, numdirauths, numrelays, numclients):
        self.total_size = total_numrelays
        self.bw_parser = bw_parser

        # 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):
        bw = int(self.calculate_relay_bandwidth())
        new_relay = relay.Relay(self.dirauthaddrs, bw, 0)
        self.relays.append(new_relay)

    def calculate_relay_bandwidth(self):
        return random.choice(bw_parser.get_distribution())

    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)

    bandwidth_file = os.getenv('BW_FILE')

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

    bw_parser = None
    if os.getenv('BW_ALGO') != "komlo":
        bw_parser = JansenBandwidthParser(bw_file=bandwidth_file, sample_size=numrelays)
    else:
        # keep the original assumption
        bw_parser = KomloBandwidthParser(sample_size=numrelays)

    total_size = bw_parser.get_relay_num()

    # 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(total_size, bw_parser, 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])