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

# import python modules
import sys
#sys.path.append("/shared/Earth/easybuild/rocky/8.4/x86_64/modules/all")
#print(sys.path)
#/shared/Earth/easybuild/rocky/8.4/x86_64/modules/all


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

from download_configurations import operational_config as dlconfig # import operational configuration
#from download_configurations import test_config as dlconfig # import test run configuration


def download_network(network, dl_data, dl_metadata, download_params, version, mode, start, stop, dl_root, meta_basefile, meta_history, q):

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

    # 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 file. '''

    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__":

    # import download configuration
    mode = dlconfig.mode
    version = dlconfig.version
    download_params = dlconfig.download_params
    
    if len(sys.argv) > 1: # take arguments from command
        path_logfiles = sys.argv[1]+'/log_files/' # string
        root = sys.argv[2] # string

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

    else: # take arguments from config file
        networks_to_download = dlconfig.networks_to_download
        path_logfiles = dlconfig.path_logfiles
        root = dlconfig.root
  

    # 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 = dlconfig.all_networks
        elif mode == 'nrt':
            networks = dlconfig.all_nrt_networks
    else:
        networks = networks_to_download


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

    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]

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

    #print(networks_to_redownload)
    #networks_to_redownload = [] # for debugging

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

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

        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]


        if (len(networks_to_reredownload) > 0) or not (networks_to_reredownload == [None]):
            logfile = '{}/log_{}_reredownload.out'.format(path_logfiles, logfile_stamp)
            print('redownload', networks_to_redownload)

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

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

        else:
            pass

        # # 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_re, 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_re)))
            
        #     # Signal the writer process to finish
        #     q_re.put('DONE')
            
        #     # Wait for the writer process to finish
        #     writer_process.join()
    
    else:
        pass



    # result = pool.starmap(download_network, zip(networks, itertools.repeat(dl_data), itertools.repeat(dl_metadata), itertools.repeat(download_params), itertools.repeat(version), itertools.repeat(mode)))
    # print(result)
    # dl_logfile.write(''.join(result))

    # pool.close()
    # pool.join()

    


    # run the desired networks
    # for network in networks:

    #     dl_start_time = time.time()
    #     max_time_per_dl = download_params[network]['max_time_dl']
    #     n_max_tries = download_params[network]['max_number_tries']
    #     dl_module = importlib.import_module(network+'_download') # import python script

    #     if dl_data == True:
    #         print("Started downloading data of {}".format(network))
    #         try:
    #             dl_module.download_data(mode, version, n_max_tries, max_time_per_dl)

    #             # copy data to gpfs
    #             #os.system('rsync -arvzh /esarchive/obs/ghost/{}/ /gpfs/projects/bsc32/AC_cache/obs/ghost/{} > rsync_gpfs.txt'.format(network_name, network_name))

    #             dl_end_time = time.time()
    #             total_dl_time = (dl_end_time - dl_start_time)/60.
    #             #dl_logfile.write('{} successful -- duration: {} minutes'.format(network_name, total_dl_time))
    #             #dl_logfile.write('\n---------------------------------------------------------------------------------------------------\n')
                

    #         except Exception as e:
    #             print(e)
    #             dl_end_time = time.time()
    #             total_dl_time = np.round((dl_end_time - dl_start_time)/60., 3)
    #             print("failed {}".format(total_dl_time))
    #             #dl_logfile.write('{} unsuccessful -- duration {} minutes \n'.format(network_name, total_dl_time)+str(traceback.format_exc()))
    #             #dl_logfile.write('\n---------------------------------------------------------------------------------------------------\n')


    #     if dl_metadata == True:
    #         print("Started downloading metadata of {}".format(network))
    #         dl_module.download_metadata(n_max_tries, max_time_per_dl)

