#!/usr/bin/env python

"""
Allows querying variables from env_*xml files and listing all available variables.

There are two usage modes:

1) Querying variables:

   - You can query a variable, or a list of variables via
      ./xmlquery var1

     or, for multiple variables (either comma or space separated)
      ./xmlquery var1,var2,var3 ....
      ./xmlquery var1 var2 var3 ....
     where var1, var2 and var3 are variables that appear in a CIME case xml file

     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
     As examples:
     - to show the number of tasks for each component, issue
        ./xmlquery NTASKS
     - to show the number of tasks just for the atm component, issue
        ./xmlquery NTASKS_ATM

     - 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 query the variable's
       value for a particular group.

       As an example, in env_batch.xml, the xml variable JOB_QUEUE appears 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 query the variable JOB_QUEUE only for one group in case.run, you need
       to specify a sub-group argument to xmlquery.
          ./xmlquery JOB_QUEUE --subgroup case.run
              JOB_QUEUE: regular
          ./xmlquery JOB_QUEUE
            Results in group case.run
                 JOB_QUEUE: regular
            Results in group case.st_archive
                 JOB_QUEUE: caldera
            Results in group case.test
                JOB_QUEUE: regular

   - You can tailor the query by adding ONE of the following possible qualifier arguments:
       [--full --fileonly --value --raw --description --get-group --type --valid-values ]
       as examples:
          ./xmlquery var1,var2 --full
          ./xmlquery var1,var2 --fileonly

   - You can query variables via a partial-match, using --partial-match or -p
       as examples:
          ./xmlquery STOP --partial-match
              Results in group run_begin_stop_restart
                  STOP_DATE: -999
                  STOP_N: 5
                  STOP_OPTION: ndays
          ./xmlquery STOP_N
                  STOP_N: 5

    - By default variable values are resolved prior to output. If you want to see the unresolved
      value(s), use the --no-resolve qualifier
      as examples:
         ./xmlquery RUNDIR
             RUNDIR: /glade/scratch/mvertens/atest/run
         ./xmlquery RUNDIR --no-resolve
             RUNDIR: $CIME_OUTPUT_ROOT/$CASE/run

2) Listing all groups and variables in those groups

      ./xmlquery --listall

     - You can list a subset of variables by adding one of the following qualifier arguments:
       [--subgroup GROUP --file FILE]

       As examples:

       If you want to see the all of the variables in group 'case.run' issue
          ./xmlquery --listall --subgroup case.run

       If you want to see all of the variables in 'env_run.xml' issue
          ./xmlquery --listall --file env_run.xml

       If you want to see all of the variables in LockedFiles/env_build.xml issue
          ./xmlquery --listall --file LockedFiles/env_build.xml

     - You can tailor the query by adding ONE of the following possible qualifier arguments:
       [--full --fileonly --raw --description --get-group --type --valid-values]

     - The env_mach_specific.xml and env_archive.xml files are not supported by this tool.
"""

from standard_script_setup import *

from CIME.case import Case
from CIME.utils import expect, convert_to_string

import textwrap, sys, re

logger = logging.getLogger("xmlquery")
unsupported_files = ["env_mach_specific.xml", "env_archive.xml"]
###############################################################################
def parse_command_line(args, description):
###############################################################################
    parser = argparse.ArgumentParser(
        description=description,
        formatter_class=argparse.RawTextHelpFormatter)

    CIME.utils.setup_standard_logging_options(parser)

    # Set command line options
    parser.add_argument("variables", nargs="*"  ,
                        help="Variable name(s) to query from env_*.xml file(s)\n"
                        "( 'variable_name' from <entry_id id='variable_name'>value</entry_id> ).\n"
                        "Multiple variables can be given, separated by commas or spaces.\n")

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

    parser.add_argument("--listall", "-listall" , default=False , action="store_true" ,
                        help="List all variables and their values.")

    parser.add_argument("--file" , "-file",
                        help="The file you want to query. If not given, queries all files.\n"
                        "Typically used with the --listall option.")

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

    parser.add_argument("-p", "--partial-match", action="store_true",
                        help="Allow partial matches of variable names, treats args as regex.")

    parser.add_argument("--no-resolve", "-no-resolve", action="store_true",
                        help="Do not resolve variable values.")

    group = parser.add_mutually_exclusive_group()

    group.add_argument("--full", default=False, action="store_true",
                       help="Print a full listing for each variable, including value, type,\n"
                       "valid values, description and file.")

    group.add_argument("--fileonly", "-fileonly", default=False, action="store_true",
                       help="Only print the filename that each variable is defined in.")

    group.add_argument("--value", "-value", default=False, action="store_true",
                       help="Only print one value without newline character.\n"
                       "If more than one has been found print first value in list.")

    group.add_argument("--raw", default=False, action="store_true",
                       help="Print the complete raw record associated with each variable.")

    group.add_argument("--description", default=False, action="store_true",
                       help="Print the description associated with each variable.")

    group.add_argument("--get-group", default=False, action="store_true",
                       help="Print the group associated with each variable.")

    group.add_argument("--type", default=False, action="store_true",
                       help="Print the data type associated with each variable.")

    group.add_argument("--valid-values", default=False, action="store_true",
                       help="Print the valid values associated with each variable, if defined.")

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

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

    if len(args.variables) == 1:
        variables = args.variables[0].split(',')
    else:
        variables = args.variables

    return variables, args.subgroup, args.caseroot, args.listall, args.fileonly, \
        args.value, args.no_resolve, args.raw, args.description, args.get_group, args.full, \
        args.type, args.valid_values, args.partial_match, args.file

def get_value_as_string(case, var, attribute=None, resolved=False, subgroup=None):
    if var in ["THREAD_COUNT", "TOTAL_TASKS", "TASKS_PER_NODE", "NUM_NODES", "SPARE_NODES", "TASKS_PER_NUMA", "CORES_PER_TASK"]:
        value = str(getattr(case, var.lower()))
    else:
        thistype = case.get_type_info(var)
        value = case.get_value(var, attribute=attribute, resolved=resolved, subgroup=subgroup)
        if value is not None and thistype:
            value = convert_to_string(value, thistype, var)

    return value

def xmlquery_sub(case, variables, subgroup=None, fileonly=False,
                 resolved=True, raw=False, description=False, get_group=False,
                 full=False, dtype=False, valid_values=False, xmlfile=None):
    """
    Return list of attributes and their values, print formatted

    """
    results = {}
    comp_classes = case.get_values("COMP_CLASSES")
    if xmlfile:
        case.set_file(xmlfile)

    # Loop over variables
    for var in variables:
        if subgroup is not None:
            groups = [subgroup]
        else:
            groups = case.get_record_fields(var, "group")
            if not groups:
                groups = ['none']

        if xmlfile:
            expect(xmlfile not in unsupported_files,
                   "XML file {} is unsupported by this tool."
                   .format(xmlfile))

            if not groups:
                value =  case.get_value(var, resolved=resolved)
                results['none'] = {}
                results['none'][var] = {}
                results['none'][var]['value']  = value
        elif not groups:
            results['none'] = {}
            results['none'][var] = {}

        for group in groups:
            if not group in results:
                results[group] = {}
            if not var in results[group]:
                results[group][var] = {}

            expect(group, "No group found for var {}".format(var))
            if get_group:
                results[group][var]['get_group'] = group

            value = get_value_as_string(case, var, resolved=resolved, subgroup=group)
            if value is None:
                var, comp, iscompvar = case.check_if_comp_var(var)
                if iscompvar:
                    value = []
                    for comp in comp_classes:
                        try:
                            nextval = get_value_as_string(case,var, attribute={"compclass" : comp}, resolved=resolved, subgroup=group)
                        except:
                            nextval = get_value_as_string(case,var, attribute={"compclass" : comp}, resolved=False, subgroup=group)

                        if nextval is not None:
                            value.append(comp + ":" + "{}".format(nextval))
                else:
                    value = get_value_as_string(case, var, resolved=resolved, subgroup=group)

            if value is None:
                if xmlfile:
                    expect(False, " No results found for variable {} in file {}".format(var, xmlfile))
                else:
                    expect(False, " No results found for variable {}".format(var))

            results[group][var]['value'] = value

            if raw:
                results[group][var]['raw'] = case.get_record_fields(var, "raw")
            if description or full:
                results[group][var]['desc'] = case.get_record_fields(var, "desc")
            if fileonly or full:
                results[group][var]['file'] = case.get_record_fields(var, "file")
            if dtype or full:
                results[group][var]['type'] = case.get_type_info(var)
            if valid_values or full:
                results[group][var]['valid_values'] = case.get_record_fields(var, "valid_values") #*** this is the problem ***

    return results

def _main_func(description):
    # Initialize command line parser and get command line options
    variables, subgroup, caseroot, listall,  fileonly, \
        value, no_resolve, raw, description, get_group, full, dtype, \
        valid_values, partial_match, xmlfile = parse_command_line(sys.argv, description)

    expect(xmlfile not in unsupported_files,
           "XML file {} is unsupported by this tool."
           .format(xmlfile))

    # Initialize case ; read in all xml files from caseroot
    with Case(caseroot) as case:
        if listall or partial_match:
            if xmlfile:
                case.set_file(xmlfile)
            all_variables = sorted(case.get_record_fields(None, "varid"))
            logger.debug("all_variables: {}".format(all_variables))
            if partial_match:
                all_matching_vars = []
                for variable in variables:
                    regex = re.compile(variable)
                    for all_variable in all_variables:
                        if regex.search(all_variable):
                            if subgroup is not None:
                                vargroups = case.get_record_fields(all_variable, "group")
                                if subgroup not in vargroups:
                                    continue

                            all_matching_vars.append(all_variable)

                variables = all_matching_vars
            else:
                if subgroup is not None:
                    all_matching_vars = []
                    for all_variable in all_variables:
                        vargroups = case.get_record_fields(all_variable, "group")
                        if subgroup not in vargroups:
                            continue
                        else:
                            all_matching_vars.append(all_variable)

                    variables = all_matching_vars
                else:
                    variables = all_variables
        expect(variables, "No variables found")
        results = xmlquery_sub(case, variables, subgroup, fileonly, resolved=not no_resolve,
                               raw=raw, description=description, get_group=get_group, full=full,
                               dtype=dtype, valid_values=valid_values, xmlfile=xmlfile)

    if full or description:
        wrapper=textwrap.TextWrapper()
        wrapper.subsequent_indent = "\t\t\t"
        wrapper.fix_sentence_endings = True
    for group in sorted(iter(results)):
        if (len(variables) > 1 or len(results) > 1 or full) and not get_group and not value:
            print("\nResults in group {}".format(group))
        cnt = 0
        for var in variables:
            if var in results[group]:
                if raw:
                    print(results[group][var]['raw'])
                elif get_group:
                    print("\t{}: {}".format(var, results[group][var]['get_group']))
                elif value:
                    if cnt > 0:
                        sys.stdout.write(",")
                    sys.stdout.write("{}".format(results[group][var]['value']))
                elif description:
                    if results[group][var]['desc'][0] is not None:
                        desc_text = ' '.join(results[group][var]['desc'][0].split())
                        print("\t{}: {}".format(var, wrapper.fill(desc_text)))
                elif fileonly:
                    print("\t{}: {}".format(var, results[group][var]['file']))
                elif dtype:
                    print("\t{}: {}".format(var, results[group][var]['type']))
                elif valid_values:
                    if 'valid_values' in results[group][var]:
                        print("\t{}: {}".format(var, results[group][var]["valid_values"]))
                elif full:
                    if results[group][var]['desc'][0] is not None:
                        desc_text = ' '.join(results[group][var]['desc'][0].split())
                    print("\t{}: value={}".format(var, results[group][var]['value']))
                    print("\t\ttype: {}".format(results[group][var]['type'][0]))
                    if 'valid_values' in results[group][var]:
                        print("\t\tvalid_values: {}".format(results[group][var]["valid_values"]))
                    print("\t\tdescription: {}".format(wrapper.fill(desc_text)))
                    print("\t\tfile: {}".format(results[group][var]['file'][0]))
                else:
                    print("\t{}: {}".format(var, results[group][var]['value']))
            cnt += 1

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