from datetime import datetime
import os
import sys
from bgpdumpy import BGPDump, TableDumpV2
import pickle
from collections import defaultdict
from radix import Radix

BGPDATA = {}
DNSDATA = ''
TRANCODATA = ''

raw_ribs = defaultdict(Radix)


def download_files(year, month, day):
    for rrc, mrtfile in BGPDATA.items():
        if not os.path.exists(mrtfile):
            remote_file = f"https://data.ris.ripe.net/{rrc}/{year}.{month:02d}/bview.{year}{month:02d}{day:02d}.0800.gz"
            os.system(f'wget {remote_file} -O {mrtfile}')


def read_bgp_data(fname):

    if not os.path.exists(fname):
        print("Warning: file does not exists {fname}")

    with BGPDump(fname) as bgp:
        print(f'Reading {fname}...')
        for entry in bgp:

            # entry.body can be either be TableDumpV1 or TableDumpV2
            if not isinstance(entry.body, TableDumpV2):
                continue  # I expect an MRT v2 table dump file

            # get a string representation of this prefix
            prefix = '%s/%d' % (entry.body.prefix, entry.body.prefixLength)
            
            if ':' in prefix:
                continue

            for route in entry.body.routeEntries:

                if not route.attr.asPath:
                    print('NO PATH?')
                    print(prefix)
                    continue

                # convert AS paths to country codes
                aspath = route.attr.asPath.split()

                peer = int(aspath[0])

                if peer not in [3257, 2914, 1299, 174, 6939]:
                    continue

                raw_ribs[peer].add(prefix)


if __name__ == '__main__':

    if len(sys.argv) < 3:
        print('usage: {sys.argv[0]} year month day')

    year = int(sys.argv[1])
    month = int(sys.argv[2])
    day = int(sys.argv[3])
    date = datetime(year, month, day)

    BGPDATA = {
            'rrc01': f'data/rrc01/bview.{year}{month:02d}{day:02d}.0800.gz',
            }

    download_files(year, month, day)

    cache_fname = f'cachev4_{year}{month:02d}{day:02d}.pickle'
    counts_fname = f'all_countsv4.pickle'

    all_counts = {3257: defaultdict(dict), 1299: defaultdict(dict), 2914: defaultdict(dict), 6939: defaultdict(dict)}
    if os.path.exists(counts_fname):
        all_counts = pickle.load(open(counts_fname, 'rb'))

    # Load BGP data
    if True: #not os.path.exists(cache_fname):
        for dump in BGPDATA.values():
            read_bgp_data(dump)

        for key in raw_ribs.keys():
            print(f'Found RIB for {key}')

        pickle.dump(raw_ribs, open(cache_fname, 'wb'))

    else:
        raw_ribs = pickle.load(open(cache_fname, 'rb'))

    for peer, rib in raw_ribs.items():
        all_counts[peer]['nb_prefix'][date] = len(rib.nodes())

    print(all_counts)
    pickle.dump(all_counts, open(counts_fname, 'wb'))