#!/usr/bin/env python
#
# Data Quality Checker also known as DQC
#
# Read the README.md for more detailed information
#
#       Barcelona Supercomputing Center - Centro Nacional de Supercomputacion
#       Earth Sciences Department - Computational Earth Sciences
#
#     Authors:
#       Joan Sala Calero (joan.sala@bsc.es)

import os, sys
import json
import cdsapi
from pathlib import Path

# Custom packages
from dqc_downloader.utils_download import *

def simplify_req(c, prod_type=None):

    # Get minimal mask (in size)
    if 'leadtime_hour' in c:
        c['leadtime_hour'] = c['leadtime_hour'][0]
    if 'time' in c:
        c['time'] = c['time'][0]
    if 'day' in c:
        c['day'] = c['day'][0]
    if 'month' in c:
        c['month'] = c['month'][0]
    if 'year' in c:
        c['year'] = c['year'][0]
    if 'product_type' in c:
        c['product_type'] = prod_type

    return c

def main():

    # Params
    if len(sys.argv) != 2:
        print('USAGE: python download_masks.py <credentials_file>')
        exit()
    else:
        cred_path = sys.argv[1]

    # URL and key to download
    with open(cred_path) as json_file:
        cred = json.load(json_file)

    # Directories to parse
    curr_dir = os.path.dirname(os.path.realpath(__file__))
    conf_dir = os.path.join(curr_dir, 'check_tests')
    masks_dir = '/tmp/masks'
    if not(os.path.exists(masks_dir)):
        os.makedirs(masks_dir)

    # Download API
    CDS = cdsapi.Client()

    # General configuration
    with open(os.path.join(conf_dir, 'download_options.json')) as json_file:
        conf = json.load(json_file)
        outputdir = conf['outdir']
        logdir = conf['logdir']
        cds_forms_c3s_location = Path(conf['formsdir'])

    # Gather datasets
    datasets = []
    for f in os.listdir(conf_dir):
        with open(os.path.join(conf_dir, f), 'r') as jsonf:
            try:
                conf = json.load(jsonf)
                for ds in conf:
                    datasets.append(ds['dataset'])
            except:
                print('ERROR: reading file {}'.format(f))

    # Gather masks
    mask_var = 'land_sea_mask'
    downloads = {}
    done = []
    for ds in datasets:
        constraints = get_constraints(ds, cds_forms_c3s_location, "constraints")
        for c in constraints:
            if mask_var in c['variable']:
                c['variable'] = mask_var
                if 'seasonal' in ds:
                    # Seasonal
                    outfile = '{}-{}-{}-{}.grib'.format(ds, c['originating_centre'][0], c['system'][0], mask_var)
                    downloads[outfile] = simplify_req(c)
                else:
                    # ERA5
                    if 'product_type' in c:
                        for p in c['product_type']:
                            outfile = '{}-{}-{}.grib'.format(ds, p, mask_var)
                            downloads[outfile] = simplify_req(c, prod_type=p)
                    # UERRA
                    if 'origin' in c:
                        outfile = '{}-{}-{}'.format(ds, c['origin'][0], mask_var)
                        downloads[outfile] = simplify_req(c)

    # Download unique masks
    for ds,c in downloads.items():
        if not(os.path.exists(os.path.join(masks_dir, outfile))):
            call_download(cred, outfile, c, ds, logdir)

if __name__ == "__main__":
    main()