#!/usr/bin/env python2.7
import os
import sys
import argparse
import yaml
import logging
# import numpy as np
from datetime import datetime


class YAMLAssistant():

    def __init__(self, yaml_fname, logger=None):
        self.yaml_fname = yaml_fname
        self.yaml = self.read_config(self.yaml_fname)
        if logger:
            self.logger = logger
        else:
            self.logger = logging.getLogger(__name__)
            self.logger.handlers = []
            self.logger.addHandler(logging.StreamHandler(sys.stdout))

    def read_config(self, yaml_fname):
        with open(yaml_fname, "r") as f:
            return yaml.load(f)

    def get_variable(self, var_name, alt_value=None, err=False):
        '''
        Extract a property defined in the yaml file.

        Nested configs can be extracted by separating the variable names with
        ':'. For example, in this yaml

        layer 1:
            layer 2: property

        'property' would be accessed by 'layer 1:layer 2'.
        '''
        nested_names = var_name.split(':')
        current_layer = self.yaml
        self.logger.debug('Variable name is' + var_name)
        for name in nested_names:
            self.logger.debug('name is' + name)
            if current_layer is None:
                return alt_value
            try:
                current_layer = current_layer.get(name)
            except KeyError:
                msg = 'Variable \"%s\" not present in the yaml file, starting'\
                      ' with \"%s\".' % (var_name, name)
                if err:
                    self.logger.error(msg)
                return alt_value
        if current_layer is None:
            return alt_value
        return current_layer

    def write_variable(self, var_name, var_contents):
        nested_names = var_name.split(':')
        current_layer = self.yaml
        while len(nested_names) > 1:
            name = nested_names.pop(0)
            if current_layer.get(name) is None:
                current_layer[name] = {}
                current_layer = current_layer[name]
            else:
                current_layer = current_layer.get(name)
        current_layer[nested_names.pop()] = var_contents
        with open(self.yaml_fname, "w") as f:
            yaml.dump(self.yaml, f, default_flow_style=False)

    # Definitions of validation functions

    def check_defined(self, var_name, err=False):
        var = self.get_variable(var_name)
        if var is None:
            msg = "Variable \"%s\" must be defined." % var_name
            if err:
                self.logger.error(msg)
                raise ValueError(msg)
            else:
                return False
        else:
            self.logger.debug("    Variable \"%s\" ok." % var_name)
            return True

    def check_nonempty_str(self, var_name, err=False):
        var = self.get_variable(var_name)
        if var is None or not isinstance(var, basestring) or len(var) == 0:
            msg = "Variable \"%s\" must be a string that is >=1 char long." \
                  % var_name
            if err:
                self.logger.error(msg)
                raise ValueError(msg)
            else:
                return False
        else:
            self.logger.debug("    Variable \"%s\" ok." % var_name)
            return True

    def check_file_exists(self, var_name, optional=False, err=False):
        var = self.get_variable(var_name)
        if var is None or not isinstance(var, basestring) \
           or not os.path.isfile(var):
            if optional:
                msg = "    Variable \"%s\" does not exist yet. May be " \
                      "generated by this program." % var_name
                self.logger.warning(msg)
            else:
                msg = "Variable \"%s\" must be present and contain a valid " \
                      "file name." % var_name
                self.logger.error(msg)
                if err:
                    self.logger.error(msg)
                    raise ValueError(msg)
                else:
                    return False
        else:
            self.logger.debug("    Variable \"%s\" ok." % var_name)
            return True

    def check_directory(self, var_name, makedirs=False, err=False):
        var = self.get_variable(var_name)
        if var is None or not isinstance(var, basestring) or \
           not os.path.isdir(var):
            if makedirs:
                makepath = os.path.abspath(var)
                os.makedirs(makepath)
                self.logger.debug("    Making directories: \"%s\"" % var)
            else:
                msg = "Variable \"%s\" must contain a valid directory; given "\
                      "\"%s\"" % (var_name, var)
                if err:
                    self.logger.error(msg)
                    raise ValueError(msg)
                else:
                    return False
        else:
            self.logger.debug("    Variable \"%s\" ok." % var_name)
            return True


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="""Read in a YAML file, and easily retrieve nested
                       variables and check variables for existence.""",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument(dest="config_file", help="""YAML configuration file""")
    parser.add_argument("--logfile", dest="log_file",
                        default="yaml_checker.log", type=str,
                        help="""Specify the log file for this run.""")
    parser.add_argument("--loglevel", dest="log_level",
                        default="INFO", type=str,
                        help="""Specify the log level for this run (DEBUG,
                                INFO, WARNING, etc)""")
    args = parser.parse_args()

    numeric_log_level = getattr(logging, args.log_level.upper(), None)
    if not isinstance(numeric_log_level, int):
        raise ValueError('Invalid log level: %s' % args.log_level)

    logger = logging.get_logger(__name__)
    logger.addHandler(logging.StreamHandler())
    logger.addHandler(logging.FileHandler(args.log_file))
    logger.setLevel(numeric_log_level)

    logger.info('Inputs:\n  config file: %s\n  log file: %s\n  log level: %s\n'
                % (args.config_file, args.log_file, args.log_level))
    logger.info('Run initiated at: %s' % (str(datetime.now())))

    yaml_chk = YAMLAssistant(args.config_file, logger)
    config = yaml_chk.read_config(args.config_file)
