"""
run the download of GHOST network data and metadata
"""

import sys
import os
import time
import datetime
import traceback
import numpy as np
import multiprocessing
import itertools
import importlib
import subprocess
from timeout import timeout
import download_parameters


def download_network(network, dl_data, dl_metadata, download_params, version, mode, start, stop, dl_root, meta_basefile, meta_history, q):
    """ downloads a given network in one pipeline
    """

    # set timeout per network
    @timeout(int(download_params[network]['total_timeout']*60)) # input in minutes
    def do_download(mode, version, n_max_tries, max_time_per_dl, start, stop, dl_root):
        dl_module.download_data(mode, version, n_max_tries, max_time_per_dl, start, stop, dl_root)

    # Start download
    dl_start_time = time.time()
    max_time_per_dl = download_params[network]['max_time_dl']
    n_max_tries = download_params[network]['max_number_tries']
    meta_params = download_params[network]['meta_params']

    if 'AERONET' in network: # take care of dots in module name
        network_mod = network.replace('.', '')
    else:
        network_mod = network

    dl_module = importlib.import_module(network_mod+'_download') # import python script

    if dl_data == True:
        print("Started downloading data of {}".format(network))
        try:
            do_download(mode, version, n_max_tries, max_time_per_dl, start, stop, dl_root)
            dl_mid_time = time.time()

            # copy data to gpfs
            if mode == "nrt":
                version = mode
            else:
                version = version

            if len(sys.argv) > 1: # take arguments from command
                subprocess.run([str(sys.argv[1])+'/rsync_data.sh', network, version, dl_root, 'original_files'])
            else: # take argument from config file
                subprocess.run(['./rsync_data.sh', network, version, dl_root, 'original_files'])

            # Stop download
            dl_end_time = time.time()

            dl_time = np.round((dl_mid_time - dl_start_time)/60., 3)
            rsync_time = np.round((dl_end_time - dl_mid_time)/60., 3)
            total_dl_time = np.round((dl_end_time - dl_start_time)/60., 3)
            res = '{} successful -- total: {} min +++ download: {} min +++ rsync: {} min \n-------------------------------------------------------------------------------------------------------------\n'.format(network, total_dl_time, dl_time, rsync_time)
            q.put(res)
            return None

        except:
            dl_end_time = time.time()
            total_dl_time = np.round((dl_end_time - dl_start_time)/60., 3)
            e = sys.exc_info()[0]
            print(e)
            res = '{} data unsuccessful -- run for {} min +++ \n'.format(network, total_dl_time)+str(traceback.format_exc())+'\n-------------------------------------------------------------------------------------------------------------\n'
            q.put(res)
            return network


    if dl_metadata == True:
        print("Started downloading metadata of {}".format(network))
        try:
            dl_module.download_metadata(n_max_tries, max_time_per_dl, version, dl_root, meta_basefile, meta_history, meta_params)
            dl_mid_time = time.time()

            # copy data to gpfs
            # if mode == "nrt":
            #     version = mode
            # else:
            #     version = version

            if len(sys.argv) > 1: # take arguments from command
                subprocess.run([str(sys.argv[1])+'/rsync_data.sh', network, version, dl_root, 'metadata'])
            else: # take argument from config file
                subprocess.run(['./rsync_data.sh', network, version, dl_root, 'metadata'])

            # Stop download
            dl_end_time = time.time()

            dl_time = np.round((dl_mid_time - dl_start_time)/60., 3)
            rsync_time = np.round((dl_end_time - dl_mid_time)/60., 3)
            total_dl_time = np.round((dl_end_time - dl_start_time)/60., 3)
            res = '{} metadata successful -- total: {} min +++ download: {} min +++ rsync: {} min \n-------------------------------------------------------------------------------------------------------------\n'.format(network, total_dl_time, dl_time, rsync_time)
            q.put(res)
        
        except:
            dl_end_time = time.time()
            total_dl_time = np.round((dl_end_time - dl_start_time)/60., 3)
            e = sys.exc_info()[0]
            print(e)
            res = '{} metadata unsuccessful -- run for {} minutes \n'.format(network, total_dl_time)+str(traceback.format_exc())+'\n-------------------------------------------------------------------------------------------------------------\n'
            q.put(res)


def do_parallel_download(q, root, mode, version, start_date, stop_date, logfile, networks, dlconfig, download_params):
    """ executes parallel downloads and returns a list of networks that failed downloading
    """

    # Create a pool of workers
    with multiprocessing.Pool(dlconfig.num_cpu) as pool:
        # Start a process to write results to the file
        writer_process = multiprocessing.Process(target=writer, args=(q, mode, version, start_date, stop_date, logfile))
        writer_process.start()
        
        # Use starmap to apply process_data to each set of arguments in args_list, passing the queue
        networks_to_redownload = pool.starmap(download_network, zip(networks, itertools.repeat(dlconfig.dl_data), itertools.repeat(dlconfig.dl_metadata), itertools.repeat(download_params), itertools.repeat(version), itertools.repeat(mode), itertools.repeat(start_date), itertools.repeat(stop_date), itertools.repeat(root), itertools.repeat(dlconfig.meta_basefile), itertools.repeat(dlconfig.meta_history), itertools.repeat(q)))
        
        # Signal the writer process to finish
        q.put('DONE')
        
        # Wait for the writer process to finish
        writer_process.join()

    return networks_to_redownload



def writer(q, mode, version, start, stop, logfile):
    '''listens for messages on the queue, writes to logfile
    '''

    dl_start_time = time.time()
    with open(logfile, 'w') as f:
        f.write('**************************************************************************************************************\n')        
        f.write('BEGINNING DOWNLOADING DATA {}, GHOST version {}, from {} to {}; started downloading {} \n'.format(mode, version, start.strftime('%Y-%m-%d'), stop.strftime('%Y-%m-%d'), datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')))
        f.write('**************************************************************************************************************\n')

        while True:
            message = q.get()
            if message == 'DONE':
                dl_end_time = time.time()
                dl_time_tot = np.round((dl_end_time - dl_start_time)/60., 3)
                f.write('Total duration: {} minutes'.format(dl_time_tot))
                break

            f.write(str(message) + '\n')
            f.flush()


if __name__ == "__main__":

    # read out command line ===============================================================================
    # testing mode flag
    if sys.argv[4] == 'true':
        testing = True
        from download_configurations import test_config as dlconfig
    else:
        testing = False
        from download_configurations import operational_config as dlconfig

    # logfile path
    path_logfiles = sys.argv[1]+'/log_files/' # string

    # root path
    root = sys.argv[2] # string
    
    # networks to download list
    if ',' in str(sys.argv[3]): 
        networks_to_download = str(sys.argv[3]).split(',')
    else:
        networks_to_download = [str(sys.argv[3])]


    # import download configuration ======================================================================
    mode = dlconfig.mode
    version = dlconfig.version
    download_params = download_parameters.download_params

    # time period for downloaded data
    if mode == 'all':
        start_date = datetime.date(1970, 1, 1) # the real date when the data starts depends on the network! This is just for representation in the logfile
        stop_date = datetime.date.today()
    elif mode == 'nrt':
        start_date = dlconfig.start_date
        stop_date = dlconfig.stop_date

    # networks to be downloaded
    if networks_to_download == ['all']:
        if mode == 'all':
            networks = download_parameters.all_networks
        elif mode == 'nrt':
            networks = download_parameters.all_nrt_networks
    else:
        networks = networks_to_download

    # print out info ====================================================================================
    print("Ghost version: {}".format(version))
    print("Mode: {}".format(mode))
    print("Number of CPUs used: {}".format(dlconfig.num_cpu))

    # create logfile
    os.makedirs(path_logfiles, exist_ok=True)
    logfile_stamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
    logfile = '{}/log_{}.out'.format(path_logfiles, logfile_stamp)

    # add queue for writing results safely to single file
    manager = multiprocessing.Manager()
    q = manager.Queue()

    # download
    networks_to_redownload = do_parallel_download(q, root, mode, version, start_date, stop_date, logfile, networks, dlconfig, download_params)
    networks_to_redownload = [i for i in networks_to_redownload if i is not None]

    # do not proceed with redownloads in testing mode
    if testing == True:
        networks_to_redownload = [] 

    # redownload =========================================================================================
    if (len(networks_to_redownload) > 0):
        logfile = '{}/log_{}_redownload.out'.format(path_logfiles,logfile_stamp)
        print('redownload', networks_to_redownload)

        # add queue for writing results safely to single file
        manager = multiprocessing.Manager()
        q_re = manager.Queue()

        # download
        networks_to_reredownload = do_parallel_download(q_re, root, mode, version, start_date, stop_date, logfile,  networks_to_redownload, dlconfig, download_params)
        networks_to_reredownload = [i for i in networks_to_reredownload if i is not None]

        # reredownload ===================================================================================
        if (len(networks_to_reredownload) > 0):
            logfile = '{}/log_{}_reredownload.out'.format(path_logfiles, logfile_stamp)
            print('reredownload', networks_to_redownload)

            # add queue for writing results safely to single file
            manager = multiprocessing.Manager()
            q_rere = manager.Queue()

            # download
            do_parallel_download(q_rere, root, mode, version, start_date, stop_date, logfile,  networks_to_reredownload, dlconfig, download_params)

        else:
            pass

    else:
        pass

