import time
import os

import requests
import xml.etree.ElementTree


class DatabricksApiException(Exception):
    """Exception class for Databricks API errors
    
    Args:
        status_code (int): HTTP status code
        error_code (int): Error specific code
        error_message (str): Error message

    Attributes:
        status_code (int): HTTP status code
        error_code (int): Error specific code
        error_message (str): Error message
    """
    def __init__(self, status_code, error_code, error_message):
        self.status_code = status_code
        self.error_code = error_code
        self.error_message = error_message


class Rest(object):
    """Rest class to execute REST get and post calls.
    JSON Results will automatically be converted to dicts
    """
    def _rest_error(self, status_code, error_code, message):
        """Create a stabdard REST error message
        
        Args:
            status_code (str): HTTP status code
            error_code (str): Error code
            message (str): Error message
        
        Returns:
            dict: Dictionary with status_code, error_code and message
        """
        return {"status_code": status_code, "error_code": error_code, "message": message}

    def _remove_tags(self, text):
        """Very simple routine to get rid of HTML tags of error responses
        
        Args:
            text (str): HTML error message
        
        Returns:
            str: The error message without HTML tags
        """
        try:
            result = ''.join(xml.etree.ElementTree.fromstring(text).itertext()).replace("\n\n", "\n")
        except:
            result = text
        return result

    def _json(self, response):
        """Convert a JSON response to dict.
        
        Args:
            response (object): Response object as returned by requests
        
        Returns:
            dict: A dict representing the JSON input or None in error case
        """
        try:
            return response.json()
        except Exception:
            raise DatabricksApiException(403, 3, "Invalid json message: %s" % self._remove_tags(response.text))

    def get(self, url, api_version, path, token):
        """REST GET function
        
        Args:
            url (str): URL for the request
            api_version (str): "1.2" (command execution) or "2.0" for all other API calls
            path (str): Additional path for the request
            token (str): Authentication token
        
        Returns:
            dict: A dict representing the JSON result of the call 
                  or in error case with keys status_code, error_code, message.
        """
        full_url = os.path.join(url, "api/%s" % api_version, path)
        try:
            response = requests.get(full_url, auth=("token", token))
        except Exception as ex:
            raise DatabricksApiException(0, 5, str(ex))
        if response.status_code in (200, 201):
            return self._json(response)
        else:
            raise DatabricksApiException(response.status_code, 1, self._remove_tags(response.text))

    def post(self, url, api_version, path, token, json=None, data=None, files=None):
        """REST POST function
        
        Args:
            url (str): URL for the request
            api_version (str): "1.2" (command execution) or "2.0" for all other API calls
            path (str): Additional path for the request
            token (str): Authentication token
            json (str, optional): json data to send in the body of the request. Defaults to None.
            data ([type], optional): Dictionary, list of tuples, bytes, or file-like object to send in 
                                     the body of the request. Defaults to None.
            files ([type], optional): Files to send with the request. Defaults to None.
            key (str, optional): Key to select from the converted JSON string. Defaults to None.
        
        Returns:
            dict: A dict representing the JSON result of the call 
            or in error case with keys status_code, error_code, message.
        """
        full_url = os.path.join(url, "api/%s" % api_version, path)
        try:
            response = requests.post(full_url, json=json, data=data, files=files, auth=("token", token))
        except Exception as ex:
            raise DatabricksApiException(0, 5, str(ex))
        if response.status_code in (200, 201):
            return self._json(response)
        else:
            raise DatabricksApiException(response.status_code, 2, self._remove_tags(response.text))


class Context(Rest):
    """Execution Context for Databricks

    Context is derived from ``Rest`` and uses Databricks REST API 1.2 
    
    Args:
        url (str): Databricks workspace URL
        cluster_id (str): Cluster ID
        token (str): Authentication token

    Attributes:
        url (str): Databricks workspace URL
        cluster_id (str): Cluster ID
        token (str): Authentication token
        id (str): Context ID
    """
    def __init__(self, url, cluster_id, token):
        self.token = token
        self.url = url
        self.cluster_id = cluster_id
        self.id = None

    def create(self):
        """Create an Execution Context
        
        Returns:
            str: Context ID
        """
        data = {"language": "python", "clusterId": self.cluster_id}
        response = self.post(self.url, "1.2", "contexts/create", data=data, token=self.token)
        self.id = response.get("id", None)
        if self.id is None:
            raise DatabricksApiException(403, 4, "Context ID missing")

    def status(self):
        """Get status of Execution Context
        
        Returns:
            dict: Dictionary with context id and status
        """
        path = "contexts/status?contextId=%s&clusterId=%s" % (self.id, self.cluster_id)
        return self.get(self.url, "1.2", path, token=self.token)

    def destroy(self):
        """Destroy Execution Context
        
        Returns:
            dict: Dictionary with context id
        """
        data = {"contextId": self.id, "clusterId": self.cluster_id}
        return self.post(self.url, "1.2", "contexts/destroy", data=data, token=self.token)


class Command(Rest):
    """Command execution on Databricks

    Command is derived from ``Rest`` and uses Databricks REST API 1.2 
        
    Args:
        url (str): Databricks workspace URL
        clusterId (str): Cluster ID
        token (str): Authentication token

    Attributes:
        url (str): Databricks workspace URL
        cluster_id (str): Cluster ID
        token (str): Authentication token
        context (Context): The Execution Context for the commands
    """
    def __init__(self, url, cluster_id, token):
        self.token = token
        self.url = url
        self.cluster_id = cluster_id
        self.context = Context(url, cluster_id, token)
        self.context.create()

    def execute(self, command):
        """Execute a python command under the given Execution Context
        
        Args:
            command (str): One line of python code
        
        Returns:
            str: The ID of the command
        """
        data = {"language": "python", "contextId": self.context.id, "clusterId": self.cluster_id}
        files = {"command": command}
        response = self.post(self.url, "1.2", "commands/execute", data=data, files=files, token=self.token)
        command_id = response.get("id", None)
        if command_id is None:
            raise DatabricksApiException(403, 4, "Command ID missing")
        
        finished = False
        while not finished:
            print(".", end="", flush=True)
            time.sleep(1)
            result = self.status(command_id)
            finished = result["status"] == "Finished"

        if result["results"]["resultType"] == "error":
            return (-1, result["results"]["cause"])
        else:
            return (0, result["results"]["data"])

    def status(self, command_id):
        """Check the status of the command execution
        
        Args:
            commandId (str): The command ID
        
        Returns:
            dict: Dictionary with context id and status
        """
        path = "commands/status?commandId=%s&contextId=%s&clusterId=%s" % (command_id, self.context.id, self.cluster_id)
        return self.get(self.url, "1.2", path, token=self.token)

    def close(self):
        self.context.destroy()
        

class Clusters(Rest):
    """Databricks Clusters API

    Clusters is derived from ``Rest`` and uses Databricks REST API 2.0
        
    Args:
        url (str): Databricks workspace URL
        token (str): Authentication token

    Attributes:
        url (str): Databricks workspace URL
        token (str): Authentication token
    """
    def __init__(self, url, token):
        self.url = url
        self.token = token

    def status(self, cluster_id):
        """Status of cluster
        
        Args:
            cluster_id (str): Cluster ID
        
        Returns:
            dict:  Cluster Response Dictionary (https://docs.databricks.com/api/latest/clusters.html#clustergetclusterresponse)
        """
        return self.get(self.url, "2.0", "clusters/get?cluster_id=%s" % cluster_id, token=self.token)

    def start(self, cluster_id):
        """Start a cluster on Databricks
        
        Args:
            cluster_id (str): Cluster ID of cluster to start
        
        Returns:
            dict:  Cluster Response Dictionary (https://docs.databricks.com/api/latest/clusters.html#clustergetclusterresponse)
        """
        data = {"cluster_id": cluster_id}
        return self.post(self.url, "2.0", "clusters/start", json=data, token=self.token)

    def edit(self, data):
        """Start a cluster on Databricks
        
        Args:
            cluster_id (str): Cluster ID of cluster to start
        
        Returns:
            dict:  Cluster Response Dictionary (https://docs.databricks.com/api/latest/clusters.html#clustergetclusterresponse)
        """
        return self.post(self.url, "2.0", "clusters/edit", json=data, token=self.token)


class Libraries(Rest):
    """Databricks Library API

    Libraries is derived from ``Rest`` and uses Databricks REST API 2.0
        
    Args:
        url (str): Databricks workspace URL
        token (str): Authentication token

    Attributes:
        url (str): Databricks workspace URL
        token (str): Authentication token
    """
    def __init__(self, url, token):
        self.url = url
        self.token = token

    def status(self, cluster_id):
        """Status of library installation
        
        Args:
            cluster_id (str): Cluster ID
        
        Returns:
            dict:  Library Response Dictionary (https://docs.databricks.com/api/latest/clusters.html#clustergetclusterresponse)
        """
        return self.get(self.url, "2.0", "libraries/cluster-status?cluster_id=%s" % cluster_id, token=self.token)
