#!/usr/bin/env python3

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

import network
import dirauth
import relay
import client

class Simulator:
    def __init__(self, relaytarget, clienttarget):
        self.relaytarget = relaytarget
        self.clienttarget = clienttarget

        # Some (for now) hard-coded parameters

        # The number of directory authorities
        numdirauths = 9

        # The fraction of relays that are fallback relays
        fracfallbackrelays = 0.05

        # 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(self.relaytarget):
            # 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))

        # The fallback relays are a hardcoded list of a small fraction
        # of the relays, used by clients for bootstrapping
        numfallbackrelays = int(self.relaytarget * fracfallbackrelays) + 1
        fallbackrelays = random.sample(self.relays, numfallbackrelays)
        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(clienttarget):
            self.clients.append(client.Client(self.dirauthaddrs))

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

        # 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()
        for c in self.clients: c.perfstats.reset()


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

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

    # 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)

    # Uncomment to see all the debug messages
    # logging.basicConfig(level=logging.DEBUG)
    logging.basicConfig(level=logging.INFO)

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

    # The steady-state numbers of relays and clients
    relaytarget = math.ceil(6500 * networkscale)
    clienttarget = math.ceil(2500000 * networkscale)

    # Create the simulation
    simulator = Simulator(relaytarget, clienttarget)

    network.thenetwork.nextepoch()

    maxmemmib = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024
    logging.info("%d MiB used", maxmemmib)