#!/usr/bin/env python3

import toml
import argparse
import sys
import collections
import os
import copy


def parse():
    """Parse the command line input arguments."""
    parser = argparse.ArgumentParser("Merge toml files")

    parser.add_argument('--toml', '-t', type=str, nargs="+", required=True, help="TOML files with run options")
    parser.add_argument('--output', '-o', required=True, nargs='?')

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

    args = parser.parse_args()

    return args.toml, args.output


def deep_update(source, overrides):
    """
    Update a nested dictionary or similar mapping.
    Modify ``source`` in place.
    """
    for key, value in overrides.items():
        if isinstance(value, collections.Mapping) and value:
            returned = deep_update(source.get(key, {}), value)
            # print("Returned:", returned)
            source[key] = returned
        else:
            override = overrides[key]
            # print("Override:", key, override)
            source[key] = override
    return source


def merge_toml_env(old_env, mods):
    return deep_update(old_env, mods)


def merge_toml_env_from_files(toml_files):
    merged_env = {}
    for toml_file in toml_files:
        if os.path.exists(toml_file):
            print(toml_file)
            modification = toml.load(open(toml_file, "r"))
            merged_env = merge_toml_env(merged_env, modification)
        else:
            print("WARNING: File not found " + toml_file)
    return merged_env


def flatten(d, sep="#"):

    obj = collections.OrderedDict()

    def recurse(t, parent_key=""):
        if isinstance(t, list):
            for i in range(len(t)):
                recurse(t[i], parent_key + sep + str(i) if parent_key else str(i))
        elif isinstance(t, dict):
            for k, v in t.items():
                recurse(v, parent_key + sep + k if parent_key else k)
        else:
            obj[parent_key] = t
    recurse(d)
    return obj


def get_member_settings(d, member, sep="#"):

    member_settings = {}
    settings = flatten(d)
    for setting in settings:
        # print(setting)
        keys = setting.split(sep)
        print(keys)
        if len(keys) == 1:
            print(member)
            member3 = "{:03d}".format(int(member))
            val = settings[setting]
            if type(val) is str:
                val = val.replace("@EEE@", member3)

            this_setting = {keys[0]: val}
            print("This setting", this_setting)
            member_settings = merge_toml_env(member_settings, this_setting)
        else:
            this_member = int(keys[-1])
            keys = keys[:-1]
            print(keys)
            if this_member == member:
                print("This is it")
                # print(setting, keys)

                this_setting = settings[setting]
                for key in reversed(keys):
                    this_setting = {key: this_setting}

                # print("TRYGVE this setting", this_setting)
                member_settings = merge_toml_env(member_settings, this_setting)
                print("TRYGVE member setting", member_settings)
    return member_settings


def merge_toml_env_from_file(toml_file):
    merged_env = {}
    if os.path.exists(toml_file):
        print(toml_file)
        modification = toml.load(open(toml_file, "r"))
        merged_env = merge_toml_env(merged_env, modification)
    else:
        print("WARNING: File not found " + toml_file)
    return merged_env


if __name__ == "__main__":

    my_files, my_output = parse()

    toml_settings = merge_toml_env_from_files(my_files)

    merged_settings = {}
    for toml_file in my_files:
        if os.path.exists(toml_file):
            print(toml_file)
            toml_settings = toml.load(open(toml_file, "r"))
        else:
            print("WARNING: File not found " + toml_file)

        for block in toml_settings:
            add_block = None

            # Treat EPS
            if block == "EPS":
                key = "ENSMSEL"
                if key in toml_settings and len(toml_settings[key]) > 0:
                    add_block = toml_settings[block]
            # Climate
            elif block == "CLIMATE":
                key = "SIMULATION_TYPE"
                if key in toml_settings and toml_settings[key] == "climate":
                    add_block = toml_settings[block]
            # Oper mode aka meps_version
            elif block == "OPER_MODE":
                add_block = toml_settings[block]
            # Extra research settings
            elif block == "RESEARCH_MODE":
                key = "RUNNING_MODE"
                if key in toml_settings and toml_settings[key] == "research":
                    add_block = toml_settings[block]
            # LETKF
            elif block == "LETKF":
                key = "ANAATMO"
                if key in toml_settings and toml_settings[key] == "LETKF":
                    print(key, toml_settings[key])
                    add_block = toml_settings[block]

            # print(block, add_block)
            if add_block is not None:
                for extra_block in add_block:
                    merged_settings = merge_toml_env(merged_settings, add_block[extra_block])
            else:
                if type(block) is not dict:
                    merged_settings.update({block: toml_settings[block]})
                else:
                    if block == "SURFEX":
                        merged_settings = merge_toml_env(merged_settings, toml_settings[block])

    # Total number of members, for the GRIB encoding
    if "ENSMSEL" in merged_settings:
        merged_settings.update({"NUMMBR": len(merged_settings["ENSMSEL"])})
    else:
        merged_settings.update({"NUMMBR": 1})

    # Write member settings
    members = None
    if "ENSMSEL" in merged_settings:
        members = merged_settings["ENSMSEL"]

    print(members, type(members), len(members))
    if members is not None and int(len(members)) > 0:
        if "EPS" in merged_settings and "MEMBER_SETTINGS" in merged_settings["EPS"]:

            for m in range(0, len(members)):
                toml_settings = copy.deepcopy(merged_settings)
                member_settings = toml_settings["EPS"]["MEMBER_SETTINGS"]

                member_dict = get_member_settings(member_settings, members[m])

                mbr = str(members[m])
                member3 = "{:03d}".format(int(members[m]))

                toml_settings = merge_toml_env(toml_settings, member_dict)

                toml_settings["EPS"].pop("MEMBER_SETTINGS", None)
                member_file = "config_" + member3 + ".toml"
                toml.dump(toml_settings, open(member_file, "w"))
            merged_settings["EPS"].pop("MEMBER_SETTINGS", None)

    # Write merged settigns
    toml.dump(merged_settings, open(my_output, "w"))
