#!/usr/bin/env python3
import surfex
import argparse
import sys
import json
import os
from datetime import datetime


def parse_args_set_assimilation_input(argv):
    """Parse the command line input arguments."""
    parser = argparse.ArgumentParser("Set assimilation input")

    # dtg, settings, sstfile = None, first_guess = None, perturbed_runs = None, lsmfile = None, obsfile = None,
    # check_existence = False, oi_coeffs = None, climfile = None):

    parser.add_argument('--version', action='version', version='surfex {0}'.format(surfex.__version__))
    parser.add_argument('--sst', required=False, default=None, type=str, nargs='?', help="Name of SST inout file")
    parser.add_argument('--ua_fg', required=False, default=None, type=str, nargs='?',
                        help="Name of upper air first guess")
    parser.add_argument('--sfx_fg', required=False, default=None, type=str, nargs='?',
                        help="Name of surfex first guess")
    parser.add_argument('--perts', required=False, default=None, type=str, nargs='+',
                        help="List of perturbed run filenames")
    parser.add_argument('--lsm', required=False, default=None, type=str, nargs='?',
                        help="Name of file with land-sea mask")
    parser.add_argument('--clim', required=False, default=None, type=str, nargs='?', help="Name of surfex first guess")
    parser.add_argument('--obs', required=False, default=None, type=str, nargs='?', help="Name of observational input")
    parser.add_argument('--oi_coeffs', required=False, default=None, type=str, nargs='?',
                        help="Name of file with OI coefficients")
    parser.add_argument('--check', '-c', action="store_true", help="Check if files exist")
    parser.add_argument('--indent', required=False, default=2, type=int, help="Indented output")
    parser.add_argument('--output', '-o', required=True, nargs='?')
    parser.add_argument('dtg', type=str, help="Date Time Group (DTG)")
    parser.add_argument('settings', type=str, help="JSON input file with settings")

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit()

    return parser.parse_args(argv)
    # return args.dtg, args.settings, args.sst, args.ua_fg, args.sfx_fg, args.perts, args.lsm, args.clim, args.obs, \
    #      args.oi_coeffs, args.check, args.indent, args.output


if __name__ == "__main__":

    # dtg, settings, sst, ua_fg, sfx_fg, perts, lsm, clim, obs, oi_coeffs, check, indent, output_file = parse()
    args = parse_args_set_assimilation_input(sys.argv[1:])
    dtg = args.dtg
    settings = args.settings
    sst = args.sst
    ua_fg = args.ua_fg
    sfx_fg = args.sfx_fg
    perts = args.perts
    lsm = args.lsm
    clim = args.clim
    obs = args.obs
    oi_coeffs = args.oi_coeffs
    check = args.check
    indent = args.indent
    output_file = args.output

    if os.path.exists(settings):
        settings = json.load(open(settings, "r"))
    else:
        raise FileNotFoundError("Settings not found: " + settings)
    dtg = datetime.strptime(dtg, "%Y%m%d%H")
    assim_input_json = surfex.set_assimilation_input(dtg, settings, sstfile=sst, ua_first_guess=ua_fg,
                                                     sfx_first_guess=sfx_fg, perturbed_runs=perts,
                                                     lsmfile=lsm, climfile=clim, obsfile=obs, oi_coeffs=oi_coeffs,
                                                     check_existence=check)

    json.dump(assim_input_json.data, open(output_file, "w"), indent=indent)
