#!/usr/bin/env python3

"""
Change a runtime parameter for SCREAM/atm. Run from your case
after case.setup.
"""

import argparse, sys, pathlib, os
import xml.etree.ElementTree as ET

# Add path to cime_config folder
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))), "cime_config"))
sys.path.append(os.path.dirname(os.path.realpath(__file__)))

from eamxx_buildnml_impl import check_value, is_array_type
from eamxx_buildnml import create_raw_xml_file
from atm_manip import atm_config_chg_impl, buffer_changes, reset_buffer, get_xml_nodes, parse_change
from utils import run_cmd_no_fail, expect

# Add path to cime
_CIMEROOT = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..","..","..","cime")
sys.path.append(os.path.join(_CIMEROOT, "CIME", "Tools"))
from standard_script_setup import * # pylint: disable=wildcard-import
from CIME.case import Case

class DeprecatedAllFlag(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        error_msg  = f"Error: '{option_string}' has been deprecated and should not be used.\n"
        error_msg +=  "  If you want to change ALL matches for a given var, use 'ANY::$varname=value' syntax instead."
        print(error_msg, file=sys.stderr)
        sys.exit(1)

###############################################################################
def recreate_raw_xml_file():
###############################################################################
    caseroot = os.getcwd()
    with Case(caseroot) as case:
        create_raw_xml_file(case, caseroot)

###############################################################################
def atm_config_chg(changes, reset=False, buffer_only=False):
###############################################################################
    if not buffer_only:
        expect(os.path.exists("namelist_scream.xml"),
               "No pwd/namelist_scream.xml file is present. Please run from a case dir that has been set up")
    else:
        expect(not reset, "Makes no sense for buffer_only and reset to both be on")

    if reset:
        reset_buffer()
        print("All buffered atmchanges have been removed.")
        hack_xml = run_cmd_no_fail("./xmlquery SCREAM_HACK_XML --value")
        if hack_xml == "TRUE":
            print("SCREAM_HACK_XML is on. Removing namelist_scream.xml to force regen")
            os.remove("namelist_scream.xml")

        recreate_raw_xml_file()
        return True
    else:
        expect(changes, "Missing <param>=<val> args")

    # Before applying/buffering changes, at the very least check the syntax
    for c in changes:
        # This will throw if the syntax is bad
        _, _, _ = parse_change(c)

    # If buffer_only=True, we must assume there were changes (we can't check).
    # Otherwise, we'll assume no changes, and if we find one, we'll adjust
    any_change = buffer_only
    if not buffer_only:
        with open("namelist_scream.xml", "r") as fd:
            tree = ET.parse(fd)
            root = tree.getroot()

        for change in changes:
            this_changed = atm_config_chg_impl(root, change)
            any_change |= this_changed


    if any_change:
        # NOTE: if a change is wrong (e.g., typo in param name), we are still buffering it.
        #       We have no way of checking this, unfortunately. If you get an error that is
        #       not just syntax, your best course of action is to run atmchange --reset.
        buffer_changes(changes)

        if not buffer_only:
            recreate_raw_xml_file()

    return True

###############################################################################
def parse_command_line(args, description):
###############################################################################
    parser = argparse.ArgumentParser(
        usage="""\n{0} <param>=<val> [<param>=<val>] ...
OR
{0} --help

\033[1mEXAMPLES:\033[0m
    \033[1;32m# Change param foo of group bar to 'hi'\033[0m
    > {0} bar::foo=hi

    \033[1;32m# Change param foo to 'hi' (only works if foo is unambiguous)\033[0m
    > {0} foo=hi

    \033[1;32m# Change all matches of param foo to 'hi'\033[0m
    > {0} ANY::foo=hi

    \033[1;32m# Change params foo to 'hi' and append 'there' to bar (only works if both are unambiguous)\033[0m
    > {0} foo=hi bar+=there
""".format(pathlib.Path(args[0]).name),
        description=description,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument(
        "-r", "--reset",
        default=False,
        action="store_true",
        help="Forget all previous atmchanges",
    )

    parser.add_argument(
        "-a", "--all",
        action=DeprecatedAllFlag,
        default=argparse.SUPPRESS,
        help="This syntax is deprecated and should not be used. Use ANY:: scope instead."
    )

    parser.add_argument(
        "-b", "--buffer-only",
        default=False,
        action="store_true",
        help="Only buffer the changes, don't actually do them. Useful for testmod scripts where the case is not setup yet",
    )

    parser.add_argument("changes", nargs="*", help="Values to change")

    return parser.parse_args(args[1:])

###############################################################################
def _main_func(description):
###############################################################################
    if "--test" in sys.argv:
        from doctest import testmod
        import atm_manip
        testmod()
        testmod(m=atm_manip)
    else:
        success = atm_config_chg(**vars(parse_command_line(sys.argv, description)))
        sys.exit(0 if success else 1)

###############################################################################

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