# Copyright 2018 Iguazio
#
# Licensed 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 ..model import RunObject
from ..utils import update_in, logger, get_in
from .base import MLRuntime, RunError
from os import path
from kubernetes import client, config
import yaml
import subprocess

class SparkK8s(MLRuntime):
    kind = 'sparkk8s'
    def __init__(self, run: RunObject):
        super().__init__(run)
        self.dashboard = ''

    def set_runtime(self, runtime):
        self.dashboard = runtime.pop("command", '')
        self.runtime = runtime

    def _run(self, runobj: RunObject):
       
        runtime = self.runtime
        extra_env = [{'name': 'MLRUN_EXEC_CONFIG', 'value': runobj.to_json()}]
        if self.rundb:
            extra_env.append({'name': 'MLRUN_META_DBPATH', 'value': self.rundb})

        update_in(runtime, 'spec.env', extra_env, append=True)

        uid = runobj.metadata.uid
        update_in(runtime, 'metadata.labels.mlrun/class', self.kind)
        update_in(runtime, 'metadata.labels.mlrun/uid', uid)
        if runobj.metadata.name:
            update_in(runtime, 'metadata.name', runobj.metadata.name)
            runtime.metadata.name = runobj.metadata.name
        name = get_in(runtime, 'metadata.name', '')
       
        #project = runobj.metadata.project

        #if not project or not name:
        #    raise RunError('name and a project must be specified in the run to deploy this function')

        #addr = deploy_config(runtime, self.dashboard, name, project, create_new=True)
        
        spark_k8s_app = self.sparkk8s_template()
        # IGUAZIO environment
        AUTH_SECRET = get_in(runtime, 'metadata.AUTH_SECRET', "NONE")
        FUSE_SECRET = get_in(runtime, 'metadata.FUSE_SECRET', "NONE")
        
        WAIT = get_in(runtime, 'metadata.WAIT', False)
        VERBOSE = get_in(runtime, 'metadata.VERBOSE', False)
        
             
        # Iguazio execution config
        #AUTH_SECRET
        spark_k8s_app['spec']['volumes'][3]['secret']['secretName'] = AUTH_SECRET
        spark_k8s_app['spec']['executor']['env'][2]['valueFrom']['secretKeyRef']['name'] = AUTH_SECRET
        spark_k8s_app['spec']['executor']['env'][3]['valueFrom']['secretKeyRef']['name'] = AUTH_SECRET
        spark_k8s_app['spec']['driver']['env'][2]['valueFrom']['secretKeyRef']['name'] = AUTH_SECRET
        spark_k8s_app['spec']['driver']['env'][3]['valueFrom']['secretKeyRef']['name'] = AUTH_SECRET


        #FUSE_SECRET
        spark_k8s_app['spec']['volumes'][4]['flexVolume']['secretRef']['name'] = FUSE_SECRET
        
        # Application behaviour
        spark_k8s_app['metadata']['name'] = get_in(runtime, 'metadata.APP_NAME', "APP_NAME")
        spark_k8s_app['spec']['mainApplicationFile'] = get_in(runtime, 'metadata.mainApplicationFile', 'local:///v3io/users/admin/spark_operator/PySparkTestUHUH.py')
        spark_k8s_app['spec']['restartPolicy']["onFailureRetries"] = get_in(runtime, 'metadata.onFailureRetries', 3)
        spark_k8s_app['spec']['restartPolicy']["onFailureRetryInterval"] = get_in(runtime, 'metadata.onFailureRetryInterval', 10)
        spark_k8s_app['spec']['restartPolicy']["onSubmissionFailureRetries"] = get_in(runtime, 'metadata.onSubmissionFailureRetries', 5)
        spark_k8s_app['spec']['restartPolicy']["onSubmissionFailureRetryInterval"] = get_in(runtime, 'metadata.onSubmissionFailureRetries', 20)
        
        # DRIVER
        spark_k8s_app['spec']['driver']['cores'] = get_in(runtime, 'metadata.driver.cores', .01)
        spark_k8s_app['spec']['driver']['coreLimit'] = get_in(runtime, 'metadata.driver.coreLimit', "200m")
        spark_k8s_app['spec']['driver']['memory'] = get_in(runtime, 'metadata.driver.memory', "512m")
        # uid of run (we need a unique value to modify the yaml for each run)
        env_runid = {'name' : 'runid', 'value': uid}
        spark_k8s_app['spec']['driver']['env'].append(env_runid)
        
        
        # Executors
        spark_k8s_app['spec']['executor']['cores'] = get_in(runtime, 'metadata.executor.cores', 1)
        spark_k8s_app['spec']['executor']['instances'] = get_in(runtime, 'metadata.executor.instances', 3)
        spark_k8s_app['spec']['executor']['memory'] = get_in(runtime, 'metadata.executor.memory', "512m")
        # uid of run (we need a unique value to modify the yaml for each run)
        env_runid = {'name' : 'runid', 'value': uid}
        spark_k8s_app['spec']['executor']['env'].append(env_runid)
        
        # Addtional JARS
        JARS = get_in(runtime, 'metadata.ADDITIONAL_JARS', "")
        spark_k8s_app['spec']['deps']['jars'].append(JARS)
        
        file = open("/tmp/" + spark_k8s_app['metadata']['name'] + ".yaml","w") 
        file.write(yaml.dump(spark_k8s_app))
        file.close()
        
        
        
        
        
        # Submit application
        runapp = subprocess.run(['kubectl', 'apply', '-f' , '/tmp/' + spark_k8s_app['metadata']['name'] + ".yaml"], stdout=subprocess.PIPE,stderr=subprocess.PIPE)
        #runappp = self.apply_yaml(path.join('/tmp/',spark_k8s_app['metadata']['name'] + ".yaml"))
        logger.info('Kubectl output: ' + runapp.stdout.decode('utf-8').replace(",",'') )
        logger.info('Submit exit Status: ' + str(runapp) )
        #logger.info('Submit exit Status: ' + str(runapp) )

        # Check for submision errors
        if runapp.returncode != 0:
            raise RunError('Failed to submit application: '+ runapp.stderr.decode('utf-8').replace("'",'') )
        
        # If wait is set
        if WAIT:
           import time
           RUNNING = "STARTING"
           while RUNNING not in [ "RUNNING", "COMPLETED" , "FAILED" ]:
               result = subprocess.run(['kubectl', 'get', 'sparkapplications', spark_k8s_app['metadata']['name'],'-o','json',"-o=jsonpath='{.status.applicationState.state}'"], stdout=subprocess.PIPE)
               RUNNING  = result.stdout.decode('utf-8').replace("'",'')
               time.sleep(10)
           
           logger.info('Application status:' + RUNNING )
           if RUNNING == "FAILED":
              raise RunError('Failed to execute application')

           wait_cmd =  subprocess.run(['kubectl', 'logs', spark_k8s_app['metadata']['name'] + '-driver','-f'], stdout=subprocess.PIPE)
           if VERBOSE:
              logger.info(wait_cmd.stdout.decode('utf-8'))
        
        
        result = subprocess.run(['kubectl', 'get', 'sparkapplications', spark_k8s_app['metadata']['name'],'-o','json',"-o=jsonpath='{.status.applicationState.state}'"], stdout=subprocess.PIPE)        
        
        logger.info('Application execution status:' + result.stdout.decode('utf-8').replace("'",''))
        return None
    



    def apply_yaml(self,sparkk8s_yaml):
        # Configs can be set in Configuration class directly or using helper
        # utility. If no argument provided, the config will be loaded from
        # default location.
        config.load_kube_config()

        with open(sparkk8s_yaml) as f:
            dep = yaml.safe_load(f)
            k8s_apps_v1 = client.AppsV1Api()
            resp = k8s_apps_v1.create_namespaced_deployment(body=dep, namespace="default")
        
        #print("Deployment created. status='%s'" % resp.metadata.name)
        return resp.metadata.name

    
    def sparkk8s_template(self):
        template={'apiVersion': 'sparkoperator.k8s.io/v1beta1',
 'kind': 'SparkApplication',
 'metadata': {'name': 'APP_NAME', 'namespace': 'default-tenant'},
 'spec': {'type': 'Python',
  'pythonVersion': '2',
  'mode': 'cluster',
  'image': 'urihoenig/spark-app-23:v01',
  'imagePullPolicy': 'Always',
  'mainApplicationFile': 'APP_FILE',
  'sparkVersion': '2.4.0',
  'restartPolicy': {'type': 'OnFailure',
   'onFailureRetries': 'RETRIES',
   'onFailureRetryInterval': 'RETRY_INTERVAL',
   'onSubmissionFailureRetries': 'SUBMISSION_FAIL_RETRIES',
   'onSubmissionFailureRetryInterval': 'SUBMISSION_FAIL_RETRIES_INTERVAL'},
  'deps': {'jars': ['/igz/java/libs/v3io-hcfs_2.11-2.3_b118_20190707140045.jar',
    '/igz/java/libs/v3io-spark2-streaming_2.11-2.3_b118_20190707140045.jar',
    '/igz/java/libs/v3io-spark2-object-dataframe_2.11-2.3_b118_20190707140045.jar',
    '/igz/java/libs/scala-library-2.11.12.jar'],
   'files': ['/igz/java/libs/v3io-py-2.3_b118_20190707140045.zip']},
  'monitoring': {'exposeDriverMetrics': True},
  'volumes': [{'hostPath': {'path': '/dev/shm/default-tenant', 'type': ''},
    'name': 'shm'},
   {'hostPath': {'path': '/var/run/iguazio/dayman/default-tenant', 'type': ''},
    'name': 'v3iod-comm'},
   {'configMap': {'defaultMode': 420, 'name': 'spark-operator-v3io-config'},
    'name': 'v3io-config'},
   {'name': 'v3io-auth',
    'secret': {'defaultMode': 420, 'secretName': 'AUTH_SECRET'}},
   {'name': 'v3io-fuse',
    'flexVolume': {'driver': 'v3io/fuse',
     'secretRef': {'name': 'FUSE_SECRET'}}},
   {'emptyDir': {}, 'name': 'daemon-health'}],
  'driver': {'cores': 0.1,
   'coreLimit': '200m',
   'memory': '512m',
   'labels': {'version': '2.4.0'},
   'serviceAccount': 'spark-operator-spark',
   'volumeMounts': [{'mountPath': '/etc/config/spark',
     'name': 'spark-master-config'},
    {'mountPath': '/var/run/iguazio/daemon_health', 'name': 'daemon-health'},
    {'mountPath': '/dev/shm', 'name': 'shm'},
    {'mountPath': '/var/run/iguazio/dayman', 'name': 'v3iod-comm'},
    {'mountPath': '/etc/config/v3io', 'name': 'v3io-config'},
    {'mountPath': '/igz/.igz', 'name': 'v3io-auth'},
    {'mountPath': '/v3io', 'name': 'v3io-fuse'}],
   'env': [{'name': 'IGZ_DATA_CONFIG_FILE',
     'value': '/igz/java/conf/v3io.conf'},
    {'name': 'CURRENT_NODE_IP',
     'valueFrom': {'fieldRef': {'apiVersion': 'v1',
       'fieldPath': 'status.hostIP'}}},
    {'name': 'V3IO_USERNAME',
     'valueFrom': {'secretKeyRef': {'key': 'username',
       'name': 'AUTH_SECRET'}}},
    {'name': 'V3IO_ACCESS_KEY',
     'valueFrom': {'secretKeyRef': {'key': 'accessKey',
       'name': 'AUTH_SECRET'}}}]},
  'executor': {'cores': 'CORES',
   'instances': 'INSTANCES',
   'memory': 'MEMORY',
   'labels': {'version': '2.4.0'},
   'volumeMounts': [{'mountPath': '/etc/config/spark',
     'name': 'spark-master-config'},
    {'mountPath': '/var/run/iguazio/daemon_health', 'name': 'daemon-health'},
    {'mountPath': '/dev/shm', 'name': 'shm'},
    {'mountPath': '/var/run/iguazio/dayman', 'name': 'v3iod-comm'},
    {'mountPath': '/etc/config/v3io', 'name': 'v3io-config'},
    {'mountPath': '/igz/.igz', 'name': 'v3io-auth'},
    {'mountPath': '/v3io', 'name': 'v3io-fuse'}],
   'env': [{'name': 'IGZ_DATA_CONFIG_FILE',
     'value': '/igz/java/conf/v3io.conf'},
    {'name': 'CURRENT_NODE_IP',
     'valueFrom': {'fieldRef': {'apiVersion': 'v1',
       'fieldPath': 'status.hostIP'}}},
    {'name': 'V3IO_USERNAME',
     'valueFrom': {'secretKeyRef': {'key': 'username',
       'name': 'AUTH_SECRET'}}},
    {'name': 'V3IO_ACCESS_KEY',
     'valueFrom': {'secretKeyRef': {'key': 'accessKey',
       'name': 'AUTH_SECRET'}}}]}}}
        return template
    
    
def run_operator():
    import yaml
    import json
    from mlrun import get_or_create_ctx, run_start


    context = get_or_create_ctx('sparkk8s')
    sparkk8s_spec = context.get_param('sparkk8sspec', {})


    # IGUAZIO environment
    AUTH_SECRET = sparkk8s_spec.get('metadata.AUTH_SECRET',"shell-um81v5npue-qh8oc-v3io-auth")
    FUSE_SECRET =  sparkk8s_spec.get('metadata.FUSE_SECRET',"shell-um81v5npue-qh8oc-v3io-fuse")

    rundb=sparkk8s_spec.get('rundb','')
    run_start(runtime=sparkk8s_spec, rundb=rundb, command='sparkk8s://')
    
if __name__ == '__main__':
        run_operator()