#!/usr/bin/env python

'''
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
'''

from PythonExecutor import PythonExecutor

import imp
import sys
import os
import pprint
import logging
import copy

logger = logging.getLogger()

class PythonReflectiveExecutor(PythonExecutor):
  """
  Some commands like STATUS, SECURITY_STATUS commands are run a lot, and need to be run really fast.
  Otherwise agent will hang waiting for them to complete every X seconds.
  
  Running the commands not in new proccess, but reflectively makes this really fast.
  """
  
  def __init__(self, tmpDir, config):
    super(PythonReflectiveExecutor, self).__init__(tmpDir, config)
    
  def run_file(self, script, script_params, tmpoutfile, tmperrfile,
               timeout, tmpstructedoutfile, callback, task_id,
               override_output_files = True, handle = None, log_info_on_failure=True):   
    pythonCommand = self.python_command(script, script_params)
    logger.debug("Running command reflectively " + pprint.pformat(pythonCommand))
    
    script_dir = os.path.dirname(script)
    self.open_subprocess_files(tmpoutfile, tmperrfile, override_output_files)
    returncode = 1

    try:
      with PythonContext(script_dir, pythonCommand):
        imp.load_source('__main__', script)
    except SystemExit as e:
      returncode = e.code
      if returncode:
        logger.debug("Reflective command failed with return_code=" + str(e))
    except Exception: 
      logger.debug("Reflective command failed with exception:", exc_info=1)
    else: 
      returncode = 0
      
    return self.prepare_process_result(returncode, tmpoutfile, tmperrfile, tmpstructedoutfile, timeout=timeout)
  
class PythonContext:
  """
  Sets and resets some context like imports, pythonpath, args.
  Also it disable logging into ambari-agent.log for reflectively called scripts.
  """
  def __init__(self, script_dir, pythonCommand):
    self.script_dir = script_dir
    self.pythonCommand = pythonCommand
    
  def __enter__(self):
    self.old_sys_path = copy.copy(sys.path)
    self.old_agv = copy.copy(sys.argv)
    self.old_sys_modules = copy.copy(sys.modules)
    self.old_logging_disable = logging.root.manager.disable
    
    logging.disable(logging.ERROR)
    sys.path.append(self.script_dir)
    sys.argv = self.pythonCommand[1:]

  def __exit__(self, exc_type, exc_val, exc_tb):
    sys.path = self.old_sys_path
    sys.argv = self.old_agv
    logging.disable(self.old_logging_disable)
    self.revert_sys_modules(self.old_sys_modules)
    return False
  
  def revert_sys_modules(self, value):
    sys.modules.update(value)
    
    for k in copy.copy(sys.modules):
      if not k in value:
        del sys.modules[k]