import base64
import glob
import logging
import os
import shutil
from typing import Callable, Dict, List

from databricks_api import DatabricksAPI
from requests import HTTPError

import forge_template.handler.handler_utils as print_utils
from forge_template.handler.handler import ACCESS_DENIED_CODE, BAD_REQUEST_CODE, NOT_FOUND_CODE, BaseHandler
from forge_template.template_exceptions import AccessDeniedException
from forge_template.template_paths import DATABRICKS_PATH, DATABRICKS_TEMPLATE_PATH, notebooks_to_copy
from forge_template.utils import YAML_CONFIG_PATH

CONFIG_PATH_PLACEHOLDER = "YAML_CONFIG_PATH_PLACEHOLDER"
DEFAULT_SCOPE_ACCESS = "WRITE"


class DatabricksHandler(BaseHandler):
    def __init__(self, name: str, config: Dict) -> None:
        config = config[name]
        self.client = DatabricksAPI(host=config["host"], token=config["token"])

        self._dev_workspace = config["development"]
        self._prod_workspace = config["production"]
        self._dbfs_path = config["config_dbfs_path"]

        self._workspaces_to_create: List[str] = []
        self._groups_to_create: List[str] = []
        self._scopes_to_create: List[Dict[str, str]] = []
        self._dbfs_dirs_to_create: List[str] = []
        self._users_to_add_to_group: Dict[str, List[str]] = {}
        self._secrets_to_add_to_scope: Dict[str, List[Dict[str, str]]] = {}
        self._groups_to_add_to_scope: Dict[str, List[str]] = {}

        self._created_workspaces: List[str] = []
        self._created_groups: List[str] = []
        self._created_scopes: List[str] = []
        self._created_dbfs_dirs: List[str] = []
        self._added_users_to_group: Dict[str, List[str]] = {}
        self._added_groups_to_scope: Dict[str, List[str]] = {}

        self._dbfs_config_path = ""

        super().__init__(name=name, config=config)

    def create_preview(self) -> None:
        self.check_dbfs_dir_to_add(self._dbfs_path)
        self.check_workspaces_to_add([self._dev_workspace["workspace_path"], self._prod_workspace["workspace_path"]])
        self.check_scopes_to_add([self._prod_workspace, self._dev_workspace])
        self.check_groups_to_add([self._prod_workspace["group"], self._dev_workspace["group"]])

    def print_preview(self):
        print_utils.print_resource_to_add(self._workspaces_to_create, "Workspaces")
        print_utils.print_resource_to_add(self._groups_to_create, "Groups")
        print_utils.print_resource_to_add(
            list(map(lambda scope: scope["scope_name"], self._scopes_to_create)), "Scopes"
        )
        print_utils.print_resource_to_add(self._dbfs_dirs_to_create, "DBFS Directories")
        logging.info("")
        print_utils.print_resources_to_update(self._users_to_add_to_group, "User", "Group")
        print_utils.print_resources_to_update(self._secrets_to_add_to_scope, "Secret", "Scope")
        print_utils.print_resources_to_update(self._groups_to_add_to_scope, "Group", "Scope")

    def setup(self):
        self.create_workspaces(self._workspaces_to_create)
        self.create_dbfs_directories(self._dbfs_dirs_to_create)
        self.create_groups(self._groups_to_create)
        self.add_users_to_group(self._users_to_add_to_group)
        self.create_scopes(self._scopes_to_create)
        self.add_groups_to_scope(self._groups_to_add_to_scope)
        self.add_secrets_to_scope(self._secrets_to_add_to_scope)

    def rollback(self) -> None:
        if self._created_scopes:
            logging.info("Rolling back all scopes that have been created.")
            self.remove_scopes(self._created_scopes)
        if self._created_groups:
            logging.info("Rolling back all groups that have been added.")
            self.remove_groups(self._created_groups)
        if self._created_workspaces:
            logging.info("Rolling back all workspaces that have been created.")
            self.remove_workspace_directories(self._created_workspaces)
        if self._created_dbfs_dirs:
            logging.info("Rolling back all DBFS directories that have been added.")
            self.remove_dbfs_directories(self._created_dbfs_dirs)
        if self._added_users_to_group:
            groups_with_new_users = print_utils.get_dict_list_subtraction(
                self._added_users_to_group, self._created_groups
            )
            if groups_with_new_users:
                logging.info("Rolling back all users that have been added to groups.")
                self.remove_users_from_groups(groups_with_new_users)
        if self._added_groups_to_scope:
            scopes_with_new_groups = print_utils.get_dict_list_subtraction(
                self._added_groups_to_scope, self._created_scopes
            )
            if scopes_with_new_groups:
                logging.info("Rolling back all groups that have been added to scopes.")
                self.remove_groups_from_scope(scopes_with_new_groups)

    def delete_all_resources(self) -> None:
        self.remove_groups([self.config["production"]["group"]["name"], self.config["development"]["group"]["name"]])
        self.remove_dbfs_directories([self.config["config_dbfs_path"]])
        self.remove_scopes([self.config["production"]["scope"]["name"], self.config["development"]["scope"]["name"]])
        self.remove_workspace_directories(
            [self.config["production"]["workspace_path"], self.config["development"]["workspace_path"]]
        )

    def upload_config(self) -> None:
        self.check_dbfs_dir_to_add(self._dbfs_path)
        self.create_dbfs_directories(self._dbfs_dirs_to_create)
        self.__copy_files_to_dbfs(self._dbfs_config_path, YAML_CONFIG_PATH)

    def create_groups(self, group_names: List[str]) -> None:
        """
        Creates Databricks groups

        Args:
            group_names: Names of groups to create
        """
        for name in group_names:
            self.client.groups.create_group(name)
            self._created_groups.append(name)
            print_utils.print_success_created(entity_name=name, entity_type="Group")

    def add_users_to_group(self, group_to_users: Dict[str, List[str]]):
        """
        Adds users to a specified group

        Args:
            group_to_users: Group name to list of user names
        """
        for group_name, users in group_to_users.items():
            for name in users:
                self.client.groups.add_to_group(group_name, user_name=name)
                if group_name not in self._added_users_to_group:
                    self._added_users_to_group[group_name] = []
                self._added_users_to_group[group_name].append(name)
                print_utils.print_success_added(group_name, "Group", name, "User")

    def check_groups_to_add(self, groups: List) -> None:
        """
        Checks which Databricks groups to add, as well as which users to add to the groups

        Args:
            groups: List of groups. See config.yaml for example.
        """
        existing_groups = self.client.groups.get_groups()["group_names"]
        for group in groups:
            group_name = group["name"]
            users = group["users"]

            if group_name in existing_groups:
                print_utils.print_resource_already_exists(group_name, "Group")
                current_users = self.__get_users_in_group(group_name)
                users_to_add = list(set(users) - set(current_users))
            else:
                users_to_add = users
                self._groups_to_create.append(group_name)

            if users_to_add:
                if group_name not in self._users_to_add_to_group:
                    self._users_to_add_to_group[group_name] = []
                self._users_to_add_to_group[group_name].extend(users_to_add)

    def __get_users_in_group(self, group_name) -> List:
        response = self.client.groups.get_group_members(group_name)
        if "members" in response:
            return [user["user_name"] for user in response["members"]]
        else:
            return []

    def remove_groups(self, groups: List) -> None:
        """
        Removes groups from Databricks

        Args:
            groups: Names of groups to remove
        """
        groups = self.check_entity_existence(groups, self.__get_existing_groups)
        for group in groups:
            self.client.groups.remove_group(group)
            print_utils.print_success_deleted(group, "Group")

    def check_entity_existence(self, entities_to_check: List, getter: Callable[[], List[str]]) -> List:
        existing_entities = getter()
        return list(set(existing_entities).intersection(set(entities_to_check)))

    def __get_existing_groups(self) -> List[str]:
        groups = self.client.groups.get_groups()
        if "group_names" not in groups:
            return []
        return groups["group_names"]

    def remove_users_from_groups(self, group_to_users: Dict) -> None:
        """
        Removes users from a Databricks group

        Args:
            group_to_users: Group name to list of user names
        """
        for group, users in group_to_users.items():
            for user in users:
                self.client.groups.remove_from_group(group, user_name=user)
                print_utils.print_success_removed(user, group, "User", "Group")

    def create_scopes(self, scopes: List[Dict]) -> None:
        """
        Create Databricks scopes

        Args:
            scopes: List with scope names
        """
        for scope in scopes:
            self.client.secret.create_scope(
                scope=scope["scope_name"], initial_manage_principal=scope["initial_manage_principal"]
            )

            self._created_scopes.append(scope["scope_name"])
            print_utils.print_success_created(scope["scope_name"], "Scope")

    def add_secrets_to_scope(self, scope_to_secrets: Dict[str, List[Dict]]) -> None:
        """
        Add secrets to a specified scope

        Args:
            scope_to_secrets: Scope name to list of secrets where each secret is a dictionary with keys:
                ``"key"``
                Key/Identifier of secret (str)

                ``"string_value"``
                The actual secret value (str)
        """
        for scope, secrets in scope_to_secrets.items():
            for secret in secrets:
                self.add_secret_to_scope(scope, secret["key"], secret["string_value"])

    def add_secret_to_scope(self, scope_name: str, key: str, value: str) -> None:
        """
        Adds a secret to a scope

        Args:
            scope_name: Name of scope
            key: Name of secret
            value: Value of secret
        """
        self.client.secret.put_secret(scope_name, key, value)
        print_utils.print_success_added(scope_name, "Scope", key, "Token")

    def remove_scopes(self, scopes: List) -> None:
        """
        Removes all scopes

        Args:
            scopes: Names of scopes to remove
        """
        scopes = self.check_entity_existence(scopes, self.__get_existing_scopes)
        for scope in scopes:
            self.client.secret.delete_scope(scope)
            print_utils.print_success_deleted(scope, "Scope")

    def remove_groups_from_scope(self, scope_to_groups: Dict) -> None:
        """
        Remove group permissions to scope in Databricks

        Args:
            scope_to_groups: Name of scope to list of group names
        """
        for scope, groups in scope_to_groups.items():
            for group in groups:
                self.client.secret.delete_acl(scope, group)
                print_utils.print_success_removed(group, scope, "Group", "Scope")

    def add_groups_to_scope(self, scope_to_group: Dict[str, List[str]]) -> None:
        """
        Add groups to a specified scope.

        Args:
            scope_to_group: Scope name to list of group names
        """
        for scope, groups in scope_to_group.items():
            for group in groups:
                try:
                    self.client.secret.put_acl(scope, group, DEFAULT_SCOPE_ACCESS)
                    if scope not in self._added_groups_to_scope:
                        self._added_groups_to_scope[scope] = []
                    self._added_groups_to_scope[scope].append(group)
                    print_utils.print_success_added(scope, "Scope", group, "Group")
                except HTTPError as e:
                    if e.response.status_code == ACCESS_DENIED_CODE:
                        logging.warning(f"Account doesn't support ACL.")
                        continue
                    raise e

    def check_scopes_to_add(self, configs: List[Dict]) -> None:
        """
        Checks which Databricks scopes to add

        Args:
            configs: List of scope configs. See config.yaml and secrets.yaml for examples.
        """
        existing_scopes = self.__get_existing_scopes()
        for config in configs:
            self.check_scope_to_add(config["scope"], existing_scopes, config["group"]["name"])

    def __get_existing_scopes(self) -> List[str]:
        scopes = self.client.secret.list_scopes()
        if "scopes" not in scopes:
            return []
        return [s["name"] for s in scopes["scopes"]]

    def check_scope_to_add(self, scope_config: Dict, existing_scopes: List[str], group_name: str) -> None:
        """
        Checks whether or not to add a given scope, as well as which secrets to add to the scope

        Args:
            scope_config: Scope configuration. See config.yaml for example.
            existing_scopes: List with names of existing scopes
            group_name: Name of group that should have access to scope
        """
        scope_name = scope_config["name"]
        secrets = scope_config["secrets"]
        initial_manage_principal = (
            self.config["initial_manage_principal"]
            if (self.config["initial_manage_principal"] != "" and self.config["initial_manage_principal"] is not None)
            else None
        )
        scope_exists = False

        if scope_name in existing_scopes:
            scope_exists = True
            print_utils.print_resource_already_exists(scope_name, "Scope")
            if self.__get_scope_acl(scope_name, group_name):
                group_name = ""
        else:
            self._scopes_to_create.append(
                {"scope_name": scope_name, "initial_manage_principal": initial_manage_principal}
            )

        if group_name:
            if scope_name not in self._groups_to_add_to_scope:
                self._groups_to_add_to_scope[scope_name] = []
            self._groups_to_add_to_scope[scope_name].append(group_name)

        if "cdf-api-key" in scope_config:
            secrets.append({"key": "cdf-api-key", "string_value": scope_config["cdf-api-key"]})

        if secrets:
            existing_secrets = self.__get_secret_names(scope_name) if scope_exists else []
            for secret in secrets:
                if secret in existing_secrets:
                    print_utils.print_resource_already_exists(secret, "Secret", override=True)

            if scope_name not in self._secrets_to_add_to_scope:
                self._secrets_to_add_to_scope[scope_name] = []
            self._secrets_to_add_to_scope[scope_name].extend(secrets)

    def __get_secret_names(self, scope_name: str) -> List[str]:
        list_secrets_response = self.client.secret.list_secrets(scope_name)
        if "secrets" in list_secrets_response:
            return [secret["key"] for secret in list_secrets_response["secrets"]]
        return []

    def __get_scope_acl(self, scope_name: str, principal: str) -> Dict:
        try:
            return self.client.secret.get_acl(scope_name, principal)
        except HTTPError as e:
            if e.response.status_code == NOT_FOUND_CODE:
                return {}
            elif e.response.status_code == BAD_REQUEST_CODE:
                return {}
            elif e.response.status_code == ACCESS_DENIED_CODE:
                return {}
            raise e

    def check_workspaces_to_add(self, workspace_paths: List) -> None:
        """
        Checks if the given workspace paths already exists.

        Args:
            workspace_paths: Workspace paths to check
        """
        for workspace_path in workspace_paths:
            if self.does_path_exist(workspace_path, lambda client, path: client.workspace.list(path)):
                print_utils.print_resource_already_exists(workspace_path, "Workspace")
            else:
                self._workspaces_to_create.append(workspace_path)

    def create_workspaces(self, workspace_paths: List) -> None:
        """
        Creates Databricks workspaces and add notebook templates to the workspaces

        Args:
            workspace_paths: Paths to add workspace to in Databricks
        """
        self.__copy_folders_locally()
        self.__copy_notebooks_locally()
        for path in workspace_paths:
            self.create_workspace(path)

    def create_workspace(self, base_path: str) -> None:
        """
        Creates a Databricks workspace and add notebook templates to the workspace

        Args:
            base_path: Path to add workspace to in Databricks
        """
        self.__create_workspace_directory(base_path)
        self._created_workspaces.append(base_path)

        for folder_name in os.listdir(DATABRICKS_PATH):
            workspace_path = os.path.join(base_path, folder_name).replace("\\", "/")
            self.__create_workspace_directory(workspace_path)
        self.__copy_notebooks_to_workspace(base_path)

    def __create_workspace_directory(self, directory_path: str) -> None:
        self.client.workspace.mkdirs(directory_path)
        print_utils.print_success_created(directory_path, "Workspace Directory")

    def __copy_notebooks_to_workspace(self, workspace_path: str) -> None:
        notebook_paths = glob.glob(os.path.join(DATABRICKS_PATH, "**", "*.ipynb"))
        for notebook_path in notebook_paths:
            file_name = os.path.basename(notebook_path)
            upload_path = os.path.join(workspace_path, notebook_path[notebook_path.index(os.sep) + 1 :]).replace(
                "\\", "/"
            )
            with open(notebook_path, "rb") as f:
                content = f.read()
                encoded_content = base64.b64encode(content).decode()
                self.client.workspace.import_workspace(
                    upload_path, language="PYTHON", format="SOURCE", content=encoded_content
                )
                print_utils.print_success_added(
                    os.path.dirname(upload_path), "Workspace Directory", file_name, "Notebook"
                )

    def __copy_notebooks_locally(self) -> None:
        for src, dest in notebooks_to_copy:
            self.__copy_files_and_replace_placeholders(src, dest)

    def __copy_files_and_replace_placeholders(self, src: str, dest: str):
        with open(src, "r") as f:
            content = f.read()

        dbfs_location = f"/dbfs/{self._dbfs_config_path[1:]}/{YAML_CONFIG_PATH}"
        content = content.replace(CONFIG_PATH_PLACEHOLDER, dbfs_location)
        with open(dest, "w") as f:
            f.write(content)

    @staticmethod
    def __copy_folders_locally() -> None:
        if not os.path.exists(DATABRICKS_PATH):
            shutil.copytree(
                DATABRICKS_TEMPLATE_PATH, DATABRICKS_PATH, ignore=shutil.ignore_patterns("*template*", "__init__.py")
            )

    def remove_workspace_directories(self, directory_paths: List) -> None:
        self.remove_paths(
            directory_paths,
            lambda client, path: client.workspace.list(path),
            lambda client, path: client.workspace.delete(path, recursive=True),
            "Workspace Directory",
        )

    def check_dbfs_dir_to_add(self, dbfs_path: str) -> None:
        """
        Checks if the given dbfs path already exists. If not, it will be added to a the list of dbfs paths to add

        Args:
            dbfs_path: Path to check
        """
        if self.does_path_exist(dbfs_path, lambda client, path: client.dbfs.list(path)):
            print_utils.print_resource_already_exists(dbfs_path, "DBFS Directory")
        else:
            self._dbfs_dirs_to_create.append(dbfs_path)

        self._dbfs_config_path = dbfs_path

    def create_dbfs_directories(self, dbfs_paths: List[str]) -> None:
        """
        Creates Databricks DBFS directories

        Args:
            dbfs_paths: Paths of DBFS directories to create
        """
        for path in dbfs_paths:
            self.create_dbfs_directory(path)

    def create_dbfs_directory(self, directory_path: str) -> None:
        """
        Creates a Databricks DBFS directory and copies the config.yaml file to that directory

        Args:
            directory_path: Path of DBFS directory to create
        """
        self.client.dbfs.mkdirs(directory_path)
        self._created_dbfs_dirs.append(directory_path)
        print_utils.print_success_created(directory_path, "DBFS Directory")
        self.__copy_files_to_dbfs(directory_path, YAML_CONFIG_PATH)

    def __copy_files_to_dbfs(self, dbfs_path: str, local_path: str) -> None:
        file_name = os.path.basename(local_path)
        with open(local_path, "rb") as f:
            content = f.read()
            encoded_content = base64.b64encode(content).decode()
            self.client.dbfs.put(os.path.join(dbfs_path, file_name).replace("\\", "/"), contents=encoded_content)
            print_utils.print_success_added(dbfs_path, "DBFS Directory", file_name, "File")

    def remove_dbfs_directories(self, directory_paths: List) -> None:
        """
        Recursively removes DBFS directories and files within those directories from Databricks

        Args:
            directory_paths: Paths of DBFS directories to remove
        """
        self.remove_paths(
            directory_paths,
            lambda client, path: client.dbfs.list(path),
            lambda client, path: client.dbfs.delete(path, recursive=True),
            "DBFS Directory",
        )

    def remove_paths(
        self,
        paths: List[str],
        existence_checker: Callable[[DatabricksAPI, str], str],
        delete_method: Callable[[DatabricksAPI, str], None],
        path_type: str,
    ) -> None:
        for path in paths:
            if self.does_path_exist(path, existence_checker):
                delete_method(self.client, path)
                print_utils.print_success_deleted(path, path_type)

    def does_path_exist(self, dir_path: str, check_function: Callable) -> bool:
        """
        Check if a directory path already exists in Databricks

        Args:
            dir_path: Path of directory to check
            check_function: Function doing the checking
        """
        try:
            check_function(self.client, dir_path)
            return True
        except HTTPError as e:
            if e.response.status_code == ACCESS_DENIED_CODE:
                raise AccessDeniedException(f"Access denied. Please check your {self.capitalized_name} access token.")
            if e.response.status_code == NOT_FOUND_CODE:
                return False
            else:
                raise e
