#!/usr/bin/env python
"""read the xml files and create a shell script that
   can be sourced to set the  environment needed for standalone NCL testing
__________________________
Created on October, 2015

Author: CSEG <cseg@cgd.ucar.edu>
"""

from __future__ import print_function

import sys

# check the system python version and require 2.7.x or greater
if sys.hexversion < 0x02070000:
    print(70 * "*")
    print("ERROR: {0} requires python >= 2.7.x. ".format(sys.argv[0]))
    print("It appears that you are running python {0}".format(
            ".".join(str(x) for x in sys.version_info[0:3])))
    print(70 * "*")
    sys.exit(1)

#
# built-in modules
#
import argparse
import os
import re
import stat

try:
    import lxml.etree as etree
except:
    import xml.etree.ElementTree as etree


# get the postprocess virtualenv path from the env_postprocess.xml file
env_file = './env_postprocess.xml'
postprocess_path = ''
standalone = ''
if os.path.isfile(env_file):
    xml_tree = etree.ElementTree()
    xml_tree.parse(env_file)
    for entry_tag in xml_tree.findall('entry'):
        if entry_tag.get('id') == 'POSTPROCESS_PATH':
            postprocess_path = entry_tag.get('value')
        if entry_tag.get('id') == 'STANDALONE':
            standalone = entry_tag.get('value')
else:
    err_msg = ('pp_config ERROR: env_postprocess.xml does not exist in this directory.')
    raise OSError(err_msg)

# check if virtualenv is activated
if hasattr(sys, 'real_prefix'):
    try:
        import jinja2
    except:
        #
        # activate the virtual environment that was created by create_python_env.sh
        #
        activate_file = '{0}/cesm-env2/bin/activate_this.py'.format(postprocess_path)
        if not os.path.isfile(activate_file):
            err_msg = ('pp_config ERROR: the virtual environment in {0} does not exist.'.format(postprocess_path) \
                          + 'Please run {0}/create_python_env.sh -cimeroot [cimeroot] -machine [machine name]'.format(postprocess_path))
            raise OSError(err_msg)

        try:
            execfile(activate_file, dict(__file__=activate_file))
        except:
            raise OSError('pp_config ERROR: Unable to activate python virtualenv {0}'.format(activate_file))
else:
    #
    # activate the virtual environment that was created by create_python_env.sh
    #
    activate_file = '{0}/cesm-env2/bin/activate_this.py'.format(postprocess_path)
    if not os.path.isfile(activate_file):
             err_msg = ('pp_config ERROR: the virtual environment in {0} does not exist.'.format(postprocess_path) \
                           + 'Please run {0}/create_python_env.sh -cimeroot [cimeroot] -machine [machine name]'.format(postprocess_path))
             raise OSError(err_msg)

    try:
        execfile(activate_file, dict(__file__=activate_file))
    except:
        raise OSError('pp_config ERROR: Unable to activate python virtualenv {0}'.format(activate_file))


if sys.version_info[0] == 2:
    from ConfigParser import SafeConfigParser as config_parser
else:
    from configparser import ConfigParser as config_parser

re_val = re.compile(r'\$(\{([A-Za-z0-9_]+)\}|[A-Za-z0-9_]+)')
#
# import modules installed in the virtual environment
#
from diag_utils import diagUtilsLib

# -------------------------------------------------------------------------------
# commandline_options - parse any command line options
# -------------------------------------------------------------------------------

def commandline_options():
    """Process the command line arguments.

    """
    parser = argparse.ArgumentParser(
        description='Read the caseroot postprocessing XML files and set the local environment.')

    parser.add_argument('-component', '--component', choices=['atm','ice','lnd','ocn'], required=True, nargs=1,
                        help='select a component [atm, ice, lnd, ocn] to set the environment using the associated XML file')

    parser.add_argument('-shell', '--shell', choices=['bash','csh'], required=True, nargs=1,
                        help='select a shell [bash, csh] to create a file that can be sourced')

    parser.add_argument('-caseroot', '--caseroot', required=False, nargs=1, default='.',
                        help='fully quailfied path to case root post processing directory. Defaults to current directory.')

    parser.add_argument('-backtrace', '--backtrace', action='store_true',
                        help='show exception backtraces as extra debugging output')

    parser.add_argument('-debug', '--debug', nargs=1, required=False, type=int, default=0,
                        help='debugging verbosity level output: 0 = none, 1 = minimum, 2 = maximum. 0 is default')

    options = parser.parse_args()
        
    return options

#==============================================================================
# expand recursive function
#==============================================================================
def expand (val, src):
  """
  Recursively traverse the environment dictionary to expand all environment variables
  
  Arguments
  val (string) - value to be expanded
  src (dictionary) -  source dictionary to search and resolve param val

  Return
  expanded value
  """
  return re_val.sub(lambda m: expand(src.get(m.group(1), ''), src), val)

#==============================================================================
# readXML - read and resolve an XML input file against a source directory
#==============================================================================
def readXML(casedir, env_file_list):
  """
  for the <entry id=value> xml files in casedir/env_file_list,
  returns a dictionary output["id"]="value"

  Arguments:
  casedir (string) - case root directory
  env_file_list - list of env_*.xml files to parse in the casedir

  Return:
  output (dictionary) 
  """
  output = dict()
  for efile in env_file_list:
    env_file = '{0}/{1}'.format(casedir, efile)
    if os.path.isfile(env_file):
      xml_tree = etree.ElementTree()
      xml_tree.parse(env_file)
      for entry_tag in xml_tree.findall('entry'):
        output[entry_tag.get('id')] = entry_tag.get('value')
    else:
      err_msg = 'cesmEnvLib.readXML ERROR: {0} does not exist or cannot be read'.format(env_file)
      raise TypeError(err_msg)

  # expand nested environment variables
  for k,v in output.iteritems():
    output[k] = expand(v, output)

  # remove () in output dictionary values
  for k,v in output.iteritems():
    output[k] = re.sub('[()]', '', v)

    # wrap the value in quotes if it contains spaces or commas
    if ',' in output[k] or ' ' in output[k]:
        new_value = '"{0}"'.format(output[k])
        output[k] = new_value
    
  return output


# -------------------------------------------------------------------------------
# main
# -------------------------------------------------------------------------------
def main(options):
    """ main
    """
    envDict = dict()

    if options.caseroot:
        caseroot = options.caseroot[0]

    prefix_map = {'atm':'ATMDIAG_',
                  'ice':'ICEDIAG_',
                  'lnd':'LNDDIAG_',
                  'ocn':'OCNDIAG_'}

    prefix = prefix_map[options.component[0]]

    compEnvFile = 'env_diags_{0}.xml'.format(options.component[0])

    envFiles = ['env_postprocess.xml', compEnvFile]

    envDict = readXML(caseroot, envFiles)

    envDict = diagUtilsLib.strip_prefix(envDict, prefix)

    envDict['CASEROOT'] = options.caseroot[0]

    filename = 'set_env_{0}'.format(options.component[0])
    first = True
    if options.shell[0] == 'csh':
        filename = '{0}/{1}.csh'.format(envDict['CASEROOT'], filename)
    else:
        filename = '{0}/{1}.sh'.format(envDict['CASEROOT'], filename)
    
    with open (filename, 'w') as envFile:
        for key, value in envDict.iteritems():
            tmpvalue = None
            if 'PATH' in key:
                tmpvalue = os.getenv(key)

            if options.shell[0] == 'csh':
                if tmpvalue is not None:
                    line = 'setenv {0} {1}:{2}\n\n'.format(key, value, tmpvalue)
                else:
                    line = 'setenv {0} {1}\n\n'.format(key, value)
                if first:
                    instruct = '*** Be sure to run "source {0}"'.format(filename)
                    first = False
            else:
                if tmpvalue is not None:
                    line = 'export {0}={1}:{2}\n\n'.format(key, value, tmpvalue)
                else:
                    line = 'export {0}={1}\n\n'.format(key, value)
                if first:
                    instruct = '*** Be sure to run "source {0}"'.format(filename)
                    first = False

            envFile.write(line)

    st = os.stat(filename)
    os.chmod(filename, st.st_mode | stat.S_IEXEC)

    return '*** Successfully created {0}\n{1}'.format(filename, instruct)
    
    
#===================================

if __name__ == "__main__":
    options = commandline_options()
    try:
        status = main(options)
        sys.exit(status)
    except Exception as error:
        print(str(error))
        sys.exit(1)
