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

# import python modules
import time
import datetime
import traceback
import numpy as np
import sys
import multiprocessing
import itertools
import importlib
import subprocess
#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, q):

    # 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']

    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:
            dl_module.download_data(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

            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)

        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)


    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)
            dl_mid_time = time.time()

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

            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 writer(q, mode, version, start, stop, logfile_path):
    '''listens for messages on the queue, writes to file. '''

    dl_start_time = time.time()
    with open(logfile_path, 'w') as f:
        f.write('**************************************************************************************************************\n')        
        f.write('BEGINNING DOWNLOADING DATA {}, GHOST version {}, from {} to {} \n'.format(mode, version, start.strftime('%Y-%m-%d'), stop.strftime('%Y-%m-%d')))
        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
    networks = dlconfig.networks_to_download

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


    print("Ghost version: {}".format(version))
    print("Mode: {}".format(mode))

    # create logfile
    logfile_path = '/esarchive/scratch/rgrodofz/software/ghost/GHOST/download_scripts/log_files/log_{}.out'.format(datetime.datetime.now().strftime('%Y%m%d-%H%M%S'))

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

    # 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_path))
        writer_process.start()
        
        # Use starmap to apply process_data to each set of arguments in args_list, passing the queue
        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(dlconfig.root), itertools.repeat(q)))
        
        # Signal the writer process to finish
        q.put('DONE')
        
        # Wait for the writer process to finish
        writer_process.join()



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

