import sys,argparse,json,time
sys.path.insert(0, './classes/')

import cohesityAPI as cohesity

# Global variables that will be used across all functions
global tagid 
tagid = None

global tagName
tagName = None

# Begin Functions

def GetArgs():
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('--cluster', '-c', type=str, action='store')
    parser.add_argument('--vcenter', '-v', type=str, action='store')
    parser.add_argument('--tag', '-t', type=str, action='store')
    parser.add_argument('--group', '-g', type=str, action='store')
    parser.add_argument('--action', '-a', type=str, action='store')
    parser.add_argument('--list', '-l', type=str, action='store')
    return (parser.parse_args())

def GetTagID(sources):
    globalList = globals()

    for source in sources:
        if globalList['tagid'] is None:
            GetNode(sources)

def GetNode(record):
    globalList = globals()

    # A non 'node' element was found, check each child key to see if one contains 'kTag'
    if "protectionSource" in record:
        if "vmWareProtectionSource" in record['protectionSource']:
            if record['protectionSource']['vmWareProtectionSource']['type'] == "kTag":

               # Tag elements have been found, look for the correct tag and assign the id to the global variable
               if record['protectionSource']['vmWareProtectionSource']['name'] == globalList['tagName']:

                    # Tag was found based on the name
                    globalList['tagid'] = record['protectionSource']['id']

    # Recurse through all 'nodes' elements looking for kTag
    if "nodes" in record:
        for node in record['nodes']:

            # The tag hasn't been found, search the next record
            if globalList['tagid'] is None:
                GetNode(node)

            # The tag was found and the global id value was set, stop iterating.
            else:
                break

def UpdateProtectionGroupExcludes(group, action):
    globalList = globals()
    isExcluded = False
    tagArray = []

    if "excludeVmTagIds" in group:
        for index in group['excludeVmTagIds']:
            for tag in index:
                tagArray.append(tag)
                if tag == globalList['tagid']:
                    isExcluded = True 
  
    # Add the tag to the exclusion
    if action == 'add':
        if isExcluded == True:
            print(globalList['tagName'] + " tag is already excluded for " + group['name'])
        else:
            print("Adding exclusion tag '" + globalList['tagName'] + "' to group " + group['name'])
            tagArray.append(globalList['tagid'])

    # Remove the tag from the exclusion
    elif action == 'remove':
        if isExcluded == True:
            print("Removing exclusion tag '" + globalList['tagName'] + "' from group " + group['name'])
            tagArray.remove(globalList['tagid'])
        else:
            print(globalList['tagName'] + " tag is not excluded for " + group['name'])


    return(tagArray)

def UpdateProtectionGroupIncludes(group, action):
    globalList = globals()
    isIncluded = False

    tagArray = []

    if "vmTagIds" in group:
        for index in group['vmTagIds']:
            for tag in index:
                tagArray.append(tag)
                if tag == globalList['tagid']:
                    isIncluded = True 
  
    # Add the tag to the inclusion
    if action == 'add':
        if isIncluded == True:
            print(globalList['tagName'] + " tag is already included for " + group['name'])
        else:
            print("Adding inclusion tag '" + globalList['tagName'] + "' to group " + group['name'])
            tagArray.append(globalList['tagid'])
            
    # Remove the tag from the inclusion
    elif action == 'remove':
        if isIncluded == True:
            print("Removing inclusion tag '" + globalList['tagName'] + "' from group " + group['name'])
            tagArray.remove(globalList['tagid'])
        else:
            print(globalList['tagName'] + " tag is not included for " + group['name'])

    return(tagArray)


# Begin Main
globalList = globals()

sourceFound = False
vCenter = None
vCenterID = None
tagIds = []

args = GetArgs()

# Validate arguments & assign variables if necessary
if not args.cluster:
    sys.exit("Error: Specify a Cohesity cluster fqdn with -c parameter.")
else:
    globalList['tagName'] = args.tag

if not args.vcenter:
    sys.exit("Error: Specify a vCenter with -v parameter")

if not args.tag:
    sys.exit("Error: Specify a tag with the -t parameter.")

if not args.action:
    sys.exit("Error: Specify an action [add|remove] with -a parameter.")
elif args.action != 'add' and args.action != 'remove':
    sys.exit("Error: Use only 'add' or 'remove' with -a parameter.")

if not args.list:
    sys.exit("Error: Specify a list [include|exclude] with -l parameter.")
elif args.list != 'exclude' and args.list != 'include':
    sys.exit("Error: Use only 'include' or 'exclude' with -l parameter.")

# Connect to the Cohesity cluster
cluster = cohesity.API(args.cluster)
authToken = cluster.GetAuthToken()
cluster.UpdateHeaders(authToken['accessToken'])

# Locate the correct vCenter and capture the id, this will be useful to validate child objects
vmSources = cluster.GetFilteredRequest("/public/protectionSources", "?environments=kVMware")
for source in vmSources:
    if source['protectionSource']['name'] == args.vcenter:
        sourceFound = True
        vCenter = source
        vCenterID = source['protectionSource']['id']
        break

# The vCenter we are looking for was not in the list of kVMware environments, stop processing.
if sourceFound == False:
    sys.exit("Error: " + args.vcenter + " is not registered to " + args.cluster)
     
# Look up the tag id based on the tagName and set a global variable when found
# IMPROVEMENT TO GET AWAY FROM GLOBAL VAR
GetTagID(vCenter)
       
if globalList['tagid'] is None: 
    sys.exit("Error: The tag '" + args.tag + "' does not exist in " + vCenter['protectionSource']['name'])
     
if args.group:
    # Replace '@' symbol with '%40' for URL encoding in the REST API
    protectionGroup = args.group.replace("@","%40")
    job = cluster.GetFilteredRequest("/public/protectionJobs", "?names=" + protectionGroup)
    
    print(json.dumps(job[0], indent=4))

    if job[0]['parentSourceId'] != vCenterID:
        sys.exit(job[0]['name'] + " does not belong to source " + vCenter['protectionSource']['name'])
    else:
        if (args.list == 'exclude'):
            excludeTagIds = UpdateProtectionGroupExcludes(job[0], args.action)
            if "excludeVmTagIds" in job[0]:
                job[0]['excludeVmTagIds'][0] = excludeTagIds
            else:
                job[0].update({"excludeVmTagIds": [ excludeTagIds ]})
        else:
            includeTagIds = UpdateProtectionGroupIncludes(job[0], args.action)
            if "vmTagIds" in job[0]:
                job[0]['vmTagIds'][0] = includeTagIds
            else:
                job[0].update({"includeVmTagIds": [ includeTagIds ]})


# FIX the screw up
        # Temp code to grab start time from previous runs
        run = cluster.GetFilteredRequest("/public/protectionRuns","?numRuns=10&jobId=" + str(job[0]['id']))
 
        try:
            hour = round(run[7]['backupRun']['stats']['startTimeUsecs']/1000000,0)
            hour = time.strftime("%H", time.localtime(hour))
            minute = round(run[7]['backupRun']['stats']['startTimeUsecs']/1000000,0)
            minute = time.strftime("%M", time.localtime(minute))

            hour = int(hour)
            minute = int(minute)

        except:
            hour = 21
            minute = 00

        job[0]['startTime']['hour'] = hour
        job[0]['startTime']['minute'] = minute 
        # EndTime

        # Temp code to build indexing
        job[0].update({"indexingPolicy":{
            "disableIndexing": False,
            "allowPrefixes": [
                "/"
            ],
            "denyPrefixes": [
                "/$Recycle.Bin",
                "/Windows",
                "/ProgramData",
                "/System Volume Information",
                "/Users/*/AppData",
                "/Recovery",
                "/usr",
                "/sys",
                "/proc",
                "/lib",
                "/grub",
                "/grub2",
                "/opt/splunk",
                "/splunk",
            ]
        }})

        # Temp code to set SLA
        job[0].update({"incrementalProtectionSlaTimeMins": 480})
        job[0].update({"fullProtectionSlaTimeMins": 480})

        job[0].update({"abortInBlackoutPeriod": False})
        job[0].update({"quiesce": False})
        job[0].update({"qosType": "kBackupHDD"})
        job[0].update({"environmentParameters":{
            "vmwareParameters": { 
                "fallbackToCrashConsistent": False,
                "skipPhysicalRdmDisks": False
            }
        }})
        job[0].update({"cloudParameters":{"failoverToCloud": False}})
        job[0].update({"leverageStorageSnapshots": False})
        job[0].update({"leverageStorageSnapshotsForHyperFlex": False})
        job[0].update({"description": ""})

# End Fix


    # Take acion on group(s)
    resp = cluster.UpdateVMProtectionJob(job[0])
    print(resp.content)

else:
    vmJobs = cluster.GetFilteredRequest("/public/protectionJobs", "?environments=kVMware")
    uniqueJobs = {job['id'] : job for job in vmJobs}.values()

    for job in uniqueJobs:
        if "isPaused" in job:
            if job['isPaused']:
                continue
        # If the job is marked for deletion, skip
        if 'isDeleted' in job:
            continue

        # If the job does not belong to this vCetner, skip
        if job['parentSourceId'] != vCenterID:
            continue

        else:
            if (args.list == 'exclude'):
                excludeTagIds = UpdateProtectionGroupExcludes(job, args.action)

                if "excludeVmTagIds" in job:
                    job['excludeVmTagIds'][0] = excludeTagIds
                else:
                    job.update({"excludeVmTagIds": [ excludeTagIds ]})

            else:
                includeTagIds = UpdateProtectionGroupIncludes(job, args.action)

                if "vmTagIds" in job:
                    job['vmTagIds'][0] = includeTagIds
                else:
                    job.update({"includeVmTagIds": [ includeTagIds ]})

# FIX the screw up
            # Temp code to grab start time from previous runs
            run = cluster.GetFilteredRequest("/public/protectionRuns","?numRuns=10&jobId=" + str(job['id']))

            try:
                hour = round(run[7]['backupRun']['stats']['startTimeUsecs']/1000000,0)
                hour = time.strftime("%H", time.localtime(hour))
                minute = round(run[7]['backupRun']['stats']['startTimeUsecs']/1000000,0)
                minute = time.strftime("%M", time.localtime(minute))

                hour = int(hour)
                minute = int(minute)

            except:
                hour = 21
                minute = 00

            job['startTime']['hour'] = hour
            job['startTime']['minute'] = minute
            # EndTime

            # Temp code to build indexing
            job.update({"indexingPolicy":{
                "disableIndexing": False,
                "allowPrefixes": [
                    "/"
                ],
                "denyPrefixes": [
                    "/$Recycle.Bin",
                    "/Windows",
                    "/ProgramData",
                    "/System Volume Information",
                    "/Users/*/AppData",
                    "/Recovery",
                    "/usr",
                    "/sys",
                    "/proc",
                    "/lib",
                    "/grub",
                    "/grub2",
                    "/opt/splunk",
                    "/splunk",
                ]
            }})

            # Temp code to set SLA
            job.update({"incrementalProtectionSlaTimeMins": 480})
            job.update({"fullProtectionSlaTimeMins": 480})

            job.update({"abortInBlackoutPeriod": False})
            job.update({"quiesce": False})
            job.update({"qosType": "kBackupHDD"})
            job.update({"environmentParameters":{
                "vmwareParameters": {
                    "fallbackToCrashConsistent": False,
                    "skipPhysicalRdmDisks": False
                }
            }})
            job.update({"cloudParameters":{"failoverToCloud": False}})
            job.update({"leverageStorageSnapshots": False})
            job.update({"leverageStorageSnapshotsForHyperFlex": False})
            job.update({"description": ""})


            # Take acion on group(s)
            print("Updating: " + job['name'])
            resp = cluster.UpdateVMProtectionJob(job)
            print(resp.content)