import sys
import os
import requests
import json
import time
import argparse
import logging
from pprint import pprint

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dbxApiVersion', help='Databricks api version', default='2.0')
    parser.add_argument('--loglevel', help='Databricks api version', default='INFO')
    requiredNamed = parser.add_argument_group('required named arguments')
    requiredNamed.add_argument('--dbxWorkspace', help='Databrics workspace name', required=True)
    requiredNamed.add_argument('--dbxClusterId', help='Databrics cluster id', required=True)
    requiredNamed.add_argument('--dbxAuthToken', help='Databrics auth token', required=True)
    return parser.parse_args()

class BearerAuth(requests.auth.AuthBase):
    def __init__(self, token):
        self.token = token
    def __call__(self, r):
        r.headers["authorization"] = "Bearer " + self.token
        return r

def dbx_cluster_status_check(databricks_url, cluster_id, auth, logger, timeout_sec=300):
    #wait for cluster to start
    timeout = time.time() + timeout_sec # 2 minutes from now
    running_cluster = requests.get(f"{databricks_url}/clusters/get", params={"cluster_id": cluster_id}, auth=auth).json()
    logger.debug(running_cluster)
    while running_cluster['state'] != "RUNNING":
        time.sleep(5)
        logger.info("Waiting for cluster to get in the running state...")
        running_cluster = requests.get(f"{databricks_url}/clusters/get", params={"cluster_id": cluster_id}, auth=auth).json()
        logger.debug(running_cluster)
        if time.time() > timeout:
            raise TimeoutError
    logger.info("Cluster is in the running state.")

def dbx_library_install(databricks_url, cluster_id, auth, library_data, logger, verify=False, poll_freq = 5, level=0, max_depth=20):
    ''' Check if library is already installed, uninstall current version, restart cluster, install new version, verify installation.
    But with recursion :-)
    '''
    if level > max_depth:
        raise RuntimeError("Recursion depth limit has been reached.")
    cluster_library_status = requests.get(f"{databricks_url}/libraries/cluster-status", params={"cluster_id": cluster_id}, auth=auth).json()
    installed_libraries = cluster_library_status.get('library_statuses',[])
    library_payload = {
        "cluster_id": cluster_id,
        "libraries": [ library_data ]
    }
    for library_status in installed_libraries:
        if library_status['library'] == library_data:
            if library_status['status'] == 'INSTALLED' and not verify:
                # uninstall
                logger.info("Found installed library, uninstalling...")
                requests.post(f"{databricks_url}/libraries/uninstall", json=library_payload, auth=auth).json()
                # wait until library changes it's status to uninstalled
                dbx_library_install(databricks_url, cluster_id, auth, library_data, logger, verify=verify, poll_freq=poll_freq, level=level+1)
            elif library_status['status'] == 'INSTALLED' and verify:
                # if verify set to true
                return
            elif library_status['status'] == 'UNINSTALL_ON_RESTART':
                # Restart the cluster
                logger.info("Found library marked for uninstall, restarting the cluster")
                response = requests.post(f"{databricks_url}/clusters/restart", json={"cluster_id": cluster_id}, auth=auth).json()
                logger.debug(response)
                dbx_cluster_status_check(databricks_url, cluster_id, databricks_auth, logger)
                # Check whether library has been removed
                dbx_library_install(databricks_url, cluster_id, auth, library_data, logger, verify=verify, poll_freq=poll_freq, level=level+1)
            elif library_status['status'] == 'SKIPPED' or library_status['status'] == 'FAILED':
                logger.error("Uninstall of the libraru has failed")
                raise RuntimeError("The library deinstallation failed")
            else:
                # Wait for status to change
                logger.info("Waiting for %s library status to stabilize.." % library_status['status'])
                time.sleep(poll_freq)
                dbx_library_install(databricks_url, cluster_id, auth, library_data, logger, verify=verify, poll_freq=poll_freq, level=level+1)
    if len(installed_libraries)==0 and not verify:
        logger.info("Installing the new version of the library")
        response = requests.post(f"{databricks_url}/libraries/install", json=library_payload, auth=auth).json()
        logger.info(response)
        # Verify that installation was successful
        dbx_library_install(databricks_url, cluster_id, auth, library_data, logger, verify=True, poll_freq=poll_freq, level=level+1)
        logger.info("Library have been sucessfully installed")

if __name__ == "__main__":
    args = parse_arguments()
    FORMAT = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    logging.basicConfig(format=FORMAT, level=args.loglevel)
    logger = logging.getLogger()

    databricks_url = f"https://{args.dbxWorkspace}/api/{args.dbxApiVersion}"
    databricks_auth=BearerAuth(args.dbxAuthToken)
    request = {
        "scope": "AWS",
        "principal": "users",
        "permission": "READ"
    }
    #response = requests.get(f"{databricks_url}/secrets/acls/list", params={"scope": "AWS"}, auth=databricks_auth).json()
    response = requests.post(f"{databricks_url}/secrets/acls/put", json=request, auth=databricks_auth).json()
    pprint(response)