#!/usr/bin/env python

"""
Allows changing variables in env_*xml files via a command-line interface.

This provides two main benefits over editing the xml files by hand:
  - Settings are checked immediately for validity
  - Settings are echoed to the CaseStatus file, providing a "paper trail" of
    changes made by the user.

Examples:

   To set a single variable:
      ./xmlchange REST_N=4

   To set multiple variables at once:
      ./xmlchange REST_OPTION=ndays,REST_N=4

   Alternative syntax (no longer recommended, but supported for backwards
   compatibility; only works for a single variable at a time):
      ./xmlchange --id REST_N --val 4

   Several xml variables that have settings for each component have somewhat special treatment.
   The variables that this currently applies to are:
    NTASKS, NTHRDS, ROOTPE, PIO_TYPENAME, PIO_STRIDE, PIO_NUMTASKS
   For example, to set the number of tasks for all components to 16, use:
      ./xmlchange NTASKS=16
   To set just the number of tasks for the atm component, use:
      ./xmlchange NTASKS_ATM=16

   The CIME case xml variables are grouped together in xml elements <group></group>.
   This is done to associate together xml variables with common features.
   Most variables are only associated with one group. However, in env_batch.xml,
   there are also xml variables that are associated with each potential batch job.
   For these variables, the '--subgroup' option may be used to specify a particular
   group for which the variable's value will be adjusted.

   As an example, in env_batch.xml, the xml variables JOB_QUEUE and JOB_WALLCLOCK_TIME
   appear in each of the batch job groups (defined in config_batch.xml):
    <group id="case.run">
    <group id="case.st_archive">
    <group id="case.test">
   To set the variable JOB_WALLCLOCK_TIME only for case.run:
      ./xmlchange JOB_WALLCLOCK_TIME=0:30 --subgroup case.run
   To set the variable JOB_WALLCLOCK_TIME for all jobs:
      ./xmlchange JOB_WALLCLOCK_TIME=0:30
"""

from standard_script_setup import *

from CIME.utils import expect, convert_to_type, append_case_status
from CIME.case import Case

import re

# Set logger
logger = logging.getLogger("xmlchange")

###############################################################################
def parse_command_line(args, description):
###############################################################################
    parser = argparse.ArgumentParser(
        description=description,
        formatter_class=argparse.RawTextHelpFormatter)

    CIME.utils.setup_standard_logging_options(parser)

    parser.add_argument("listofsettings", nargs="?", default='',
                        help="Comma-separated list of settings in the form: var1=value,var2=value,...")

    parser.add_argument("--caseroot", default=os.getcwd(),
                        help="Case directory to change.\n"
                        "Default is current directory.")

    # Need to support older single dash version of arguments for compatibility with components

    parser.add_argument("--append","-append", action="store_true",
                        help="Append to the existing value rather than overwriting it.")

    parser.add_argument("--subgroup","-subgroup",
                        help="Apply to this subgroup only.")

    parser.add_argument("--id", "-id",
                        help="The variable to set.\n"
                        "(Used in the alternative --id var --val value form, rather than\n"
                        "the recommended var=value form.)")

    parser.add_argument("--val","-val",
                        help="The value to set.\n"
                        "(Used in the alternative --id var --val value form, rather than\n"
                        "the recommended var=value form.)")

    parser.add_argument("--file", "-file",
                        help="XML file to edit.\n"
                        "Generally not needed, but can be specified to ensure that only the\n"
                        "expected file is being changed. (If a variable is not found in this file,\n"
                        "an error will be generated.)")

    parser.add_argument("--delimiter","-delimiter", type=str, default="," ,
                        help="Delimiter string in listofvalues.\n"
                        "Default is ','.")

    parser.add_argument("--dryrun","-dryrun", action="store_true",
                        help="Parse settings and print key-value pairs, but don't actually change anything.")

    parser.add_argument("--noecho", "-noecho", action="store_true",
                        help="Do not update CaseStatus with this change.\n"
                        "This option is mainly meant to be used by cime scripts: the 'paper trail' in\n"
                        "CaseStatus is meant to show changes made by the user, so we generally don't\n"
                        "want this to be contaminated by changes made automatically by cime scripts.")

    parser.add_argument("-f","--force", action="store_true",
                        help="Ignore typing checks and store value.")

    parser.add_argument("-loglevel",
                        help="Ignored, only for backwards compatibility.")

    args = CIME.utils.parse_args_and_handle_standard_logging_options(args, parser)

    listofsettings = []
    if( len(args.listofsettings )):
        expect(args.id is None, "Cannot specify both listofsettings and --id")
        expect(args.val is None, "Cannot specify both listofsettings and --val")
        delimiter = re.escape(args.delimiter)
        listofsettings = re.split(r'(?<!\\)'+ delimiter , args.listofsettings)
    else:
        expect(args.id is not None and args.val is not None,
               "Must give either (1) listofsettings or (2) both --id and --val")

    return args.caseroot, listofsettings, args.file, args.id, args.val, args.subgroup, args.append, args.noecho, args.force , args.dryrun

def xmlchange_single_value(case, xmlid, xmlval, subgroup, append, force, dryrun):
    if xmlid in ["THREAD_COUNT", "TOTAL_TASKS", "TASKS_PER_NODE", "NUM_NODES", "SPARE_NODES", "TASKS_PER_NUMA", "CORES_PER_TASK"]:
        expect(False, "Cannot xmlchange derived attribute {}".format(xmlid))

    type_str = case.get_type_info(xmlid)
    if append:
        value = case.get_value(xmlid, resolved=False,
                               subgroup=subgroup)
        xmlval = "%s %s" % (value, xmlval)

    if type_str is not None and not force:
        xmlval = convert_to_type(xmlval, type_str, xmlid)

    if not dryrun:
        newval = case.set_value(xmlid, xmlval, subgroup, ignore_type=force)
        expect(newval is not None,"No variable \"%s\" found"%xmlid)
    else:
        logger.warning("'%s' = '%s'" , xmlid , xmlval )

    # Subtle: If the user is making specific requests for walltime and queue, we need to
    # store these choices in USER_REQUESTED_WALLTIME and/or USER_REQUESTED_QUEUE so they
    # are not lost if the user does a case.setup --reset.
    if xmlid == "JOB_WALLCLOCK_TIME":
        case.set_value("USER_REQUESTED_WALLTIME", xmlval, subgroup)

    if xmlid == "JOB_QUEUE":
        case.set_value("USER_REQUESTED_QUEUE", xmlval, subgroup)

def xmlchange(caseroot, listofsettings, xmlfile, xmlid, xmlval, subgroup,
              append, noecho, force, dryrun):

    with Case(caseroot, read_only=False) as case:
        comp_classes = case.get_values("COMP_CLASSES")
        if xmlfile:
            case.set_file(xmlfile)
        case.set_comp_classes(comp_classes)

        if len(listofsettings):
            logger.debug("List of attributes to change: %s" , listofsettings)

            # Change values
            for setting in listofsettings:

                pair = setting.split("=",1)
                expect(len(pair) == 2 , "Expecting a key value pair in the form of key=value. Got %s" % (pair) )
                (xmlid, xmlval) = pair

                xmlchange_single_value(case, xmlid, xmlval, subgroup, append, force, dryrun)
        else:
            xmlchange_single_value(case, xmlid, xmlval, subgroup, append, force, dryrun)

    if not noecho:
        argstr = ""
        for arg in sys.argv:
            argstr += "%s "%arg
        msg = "<command> %s </command>" % (argstr)
        append_case_status("xmlchange", "success", msg=msg, caseroot=caseroot)

def _main_func(description):
    # pylint: disable=unused-variable
    caseroot, listofsettings, xmlfile, xmlid, xmlval, subgroup, append, noecho, force , dry = parse_command_line(sys.argv, description)

    xmlchange(caseroot, listofsettings, xmlfile, xmlid, xmlval, subgroup, append, noecho, force, dry)

if (__name__ == "__main__"):
    _main_func(__doc__)