#!/usr/bin/env python3

# To run this script you need to install DatabricksAPI "pip install databricks-api"
# find more details here: https://github.com/crflynn/databricks-api
#
# REQUIRED VARIABLES
# dataBricksHost (Default: dev-delphi.cloud.databricks.com)
# dataBricksUsername
# dataBricksPassword

from databricks_api import DatabricksAPI
import re
import os, sys, getopt

re_job_cluster_name = re.compile('job-([\d]+)-run-[\d]+', re.IGNORECASE)

def suggest_tag_by_name(object_name, tag_name):
  if tag_name =='project':
    if "exploration" in object_name.lower():
        return "Exploration Area"
    return 'Apps ETL'
  if tag_name =='service':
    return 'Databricks'
  if tag_name =='plat_env_project_service':
    tag_val="DLP_"
    if "stag" in object_name.lower():
        tag_val=tag_val+"STAGE_"
    else:
        if "qa" in object_name.lower():
                tag_val=tag_val+"QA_"
        else:
            if "test" in object_name.lower():
                tag_val=tag_val+"QA_"
            else:
                tag_val=tag_val+"DEV_"
    if "exploration" in object_name.lower():
        tag_val=tag_val+"EXP_"
    else:
        tag_val=tag_val+"APPE_"
    tag_val=tag_val+"DBX"
    return tag_val
  if tag_name =='environment':
    if "stag" in object_name.lower():
        return "Staging"
    if "qa" in object_name.lower():
        return "QA"
    if "test" in object_name.lower():
        return "QA"
    return 'Development'
  if tag_name =='platform':
    return 'Delphi'
  if tag_name =='platform_with_env':
    return 'unidentified'

def tags_changed(old_tags,new_tags):
    for each_tagname in new_tags.keys():
        if each_tagname not in old_tags:
            return True
        if old_tags[each_tagname] != new_tags[each_tagname]:
            return True
    return False

def get_fixed_tags_by_name(object_name,tags,required_tag_names):
    fixedJobTags=dict(tags)
    for each_tagname in required_tag_names:
        if each_tagname not in fixedJobTags:
            fixedJobTags[each_tagname]=suggest_tag_by_name(object_name,each_tagname)
    return fixedJobTags

if __name__ == "__main__":
    required_tag_names = [
        'project',
        'service',
        'plat_env_project_service',
        'environment',
        'platform'
    ]

    if "dataBricksHost" in os.environ:
        dataBricksHost = os.environ["dataBricksHost"]
    else:
        dataBricksHost = "dev-delphi.cloud.databricks.com"
    if "dataBricksUsername" in os.environ:
        dataBricksUsername = os.environ["dataBricksUsername"]
    else:
        print("ERROR: variable dataBricksUsername is not declared in ENV")
        exit(1)
    if "dataBricksPassword" in os.environ:
        dataBricksPassword = os.environ["dataBricksPassword"]
    else:
        print("ERROR: variable dataBricksPassword is not declared in ENV")
        exit(1)

    dbricks = DatabricksAPI(
        host=dataBricksHost,
        user=dataBricksUsername,
        password=dataBricksPassword
    )

    # clusterList = dbricks.cluster.list_clusters(headers=None)
    # for each_cluster in clusterList["clusters"]:
    #     if re_job_cluster_name.findall(each_cluster["cluster_name"]):
    #         print ('JOB: Cluster:',each_cluster["cluster_name"], 'has no TAGS.')
    #     if "custom_tags" in each_cluster:
    #     #    print(each_cluster["cluster_name"])
    #     #    print(each_cluster["custom_tags"])
    #     #    print("====")
    #         print()
    #     else:
    #         print ('WARNING: Cluster:',each_cluster["cluster_name"], 'has no TAGS.')

    try:
      opts, args = getopt.getopt(sys.argv[1:],"",["apply"])
    except getopt.GetoptError:
      print(sys.argv[0],' -y - Apply changes')
      sys.exit(2)
    apply = False
    for opt, arg in opts:
      if opt == '-h':
         print(sys.argv[0],' -y - Apply changes')
         sys.exit()
      elif opt in ("-y", "--apply"):
         apply = True

    jobList = dbricks.jobs.list_jobs(headers=None)
    for each_job in jobList["jobs"]:
        fixedJobTags={}
        jobTags={}
        if "new_cluster" in each_job["settings"]:
            if "custom_tags" in each_job["settings"]["new_cluster"]:
                jobTags=each_job["settings"]["new_cluster"]["custom_tags"]
                #print(each_job["settings"]["new_cluster"]["custom_tags"]["plat_env_project_service"])
            fixedJobTags=get_fixed_tags_by_name(each_job["settings"]["name"],jobTags,required_tag_names)        
        #else:
            # "New Cluster" not configured for this Job
        if tags_changed(jobTags,fixedJobTags):
            print("\nUPDATING JOB: ",each_job["settings"]["name"], "\n==============")
            print("OLD-Tags", jobTags, "\nNEW-Tags ",fixedJobTags)
            each_job["settings"]["new_cluster"]["custom_tags"]=fixedJobTags
            if apply:
                print('Applying changes...')
                dbricks.jobs.reset_job(
                    each_job["job_id"],
                    each_job["settings"],
                    headers=None,
                )
                print('DONE')
            else:
                print('To apply changes run :',sys.argv[0],' -y')
        else:
            print('No changes to be applyed')