#
#   Copyright 2020 Logical Clocks AB
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#

import os
import boto3
import base64
import json
import requests

try:
    from pyspark.sql import SparkSession
except ImportError:
    pass

from hsfs.client import base, auth, exceptions
from hsfs.client.exceptions import FeatureStoreException


class Client(base.Client):
    DEFAULT_REGION = "default"
    SECRETS_MANAGER = "secretsmanager"
    PARAMETER_STORE = "parameterstore"
    LOCAL_STORE = "local"

    def __init__(
        self,
        host,
        port,
        project,
        engine,
        region_name,
        secrets_store,
        hostname_verification,
        trust_store_path,
        cert_folder,
        api_key_file,
        api_key_value,
    ):
        """Initializes a client in an external environment such as AWS Sagemaker."""
        if not host:
            raise exceptions.ExternalClientError("host")
        if not project:
            raise exceptions.ExternalClientError("project")

        self._host = host
        self._port = port
        self._base_url = "https://" + self._host + ":" + str(self._port)
        self._project_name = project
        self._region_name = region_name or self.DEFAULT_REGION

        if api_key_value is not None:
            api_key = api_key_value
        else:
            api_key = self._get_secret(secrets_store, "api-key", api_key_file)
        self._auth = auth.ApiKeyAuth(api_key)

        self._session = requests.session()
        self._connected = True
        self._verify = self._get_verify(self._host, trust_store_path)

        project_info = self._get_project_info(self._project_name)
        self._project_id = str(project_info["projectId"])

        self._cert_key = None
        self._cert_folder_base = None

        if engine == "python":
            # On external Spark clients (Databricks, Spark Cluster),
            # certificates need to be provided before the Spark application starts.
            self._cert_folder_base = cert_folder
            self._cert_folder = os.path.join(cert_folder, host, project)
            self._trust_store_path = os.path.join(self._cert_folder, "trustStore.jks")
            self._key_store_path = os.path.join(self._cert_folder, "keyStore.jks")

            os.makedirs(self._cert_folder, exist_ok=True)
            credentials = self._get_credentials(self._project_id)
            self._write_b64_cert_to_bytes(
                str(credentials["kStore"]),
                path=self._get_jks_key_store_path(),
            )
            self._write_b64_cert_to_bytes(
                str(credentials["tStore"]),
                path=self._get_jks_trust_store_path(),
            )

            self._write_pem_file(credentials["caChain"], self._get_ca_chain_path())
            self._write_pem_file(
                credentials["clientCert"], self._get_client_cert_path()
            )
            self._write_pem_file(credentials["clientKey"], self._get_client_key_path())

            self._cert_key = str(credentials["password"])
            with open(os.path.join(self._cert_folder, "material_passwd"), "w") as f:
                f.write(str(credentials["password"]))

        elif engine.startswith("spark"):
            _spark_session = SparkSession.builder.getOrCreate()

            self.validate_spark_configuration(
                _spark_session, engine == "spark-no-metastore"
            )
            with open(
                _spark_session.conf.get("spark.hadoop.hops.ssl.keystores.passwd.name"),
                "r",
            ) as f:
                self._cert_key = f.read()

            self._trust_store_path = _spark_session.conf.get(
                "spark.hadoop.hops.ssl.trustore.name"
            )
            self._key_store_path = _spark_session.conf.get(
                "spark.hadoop.hops.ssl.keystore.name"
            )

    def validate_spark_configuration(self, _spark_session, no_metastore):
        exception_text = "Spark is misconfigured for communication with Hopsworks, missing or invalid property: "

        configuration_dict = {
            "spark.hadoop.hops.ssl.trustore.name": None,
            "spark.hadoop.hops.rpc.socket.factory.class.default": "io.hops.hadoop.shaded.org.apache.hadoop.net.HopsSSLSocketFactory",
            "spark.serializer": "org.apache.spark.serializer.KryoSerializer",
            "spark.hadoop.hops.ssl.hostname.verifier": "ALLOW_ALL",
            "spark.hadoop.hops.ssl.keystore.name": None,
            "spark.hadoop.fs.hopsfs.impl": "io.hops.hopsfs.client.HopsFileSystem",
            "spark.hadoop.hops.ssl.keystores.passwd.name": None,
            "spark.hadoop.hops.ipc.server.ssl.enabled": "true",
            "spark.hadoop.client.rpc.ssl.enabled.protocol": "TLSv1.2",
        }

        if not no_metastore:
            configuration_dict["spark.hadoop.hive.metastore.uris"] = None
            configuration_dict["spark.sql.hive.metastore.jars"] = None

        for key, value in configuration_dict.items():
            if not (
                _spark_session.conf.get(key, "not_found") != "not_found"
                and (value is None or _spark_session.conf.get(key, None) == value)
            ):
                raise FeatureStoreException(exception_text + key)

    def _close(self):
        """Closes a client and deletes certificates."""
        if self._cert_folder_base is None:
            # On external Spark clients (Databricks, Spark Cluster),
            # certificates need to be provided before the Spark application starts.
            return

        # Clean up only on AWS
        self._cleanup_file(self._get_jks_key_store_path())
        self._cleanup_file(self._get_jks_trust_store_path())
        self._cleanup_file(os.path.join(self._cert_folder, "material_passwd"))
        self._cleanup_file(self._get_ca_chain_path())
        self._cleanup_file(self._get_client_cert_path())
        self._cleanup_file(self._get_client_key_path())

        try:
            # delete project level
            os.rmdir(self._cert_folder)
            # delete host level
            os.rmdir(os.path.dirname(self._cert_folder))
            # on AWS base dir will be empty, and can be deleted otherwise raises OSError
            os.rmdir(self._cert_folder_base)
        except OSError:
            pass
        self._connected = False

    def _get_jks_trust_store_path(self):
        return self._trust_store_path

    def _get_jks_key_store_path(self):
        return self._key_store_path

    def _get_ca_chain_path(self) -> str:
        return os.path.join(self._cert_folder, "ca_chain.pem")

    def _get_client_cert_path(self) -> str:
        return os.path.join(self._cert_folder, "client_cert.pem")

    def _get_client_key_path(self) -> str:
        return os.path.join(self._cert_folder, "client_key.pem")

    def _get_secret(self, secrets_store, secret_key=None, api_key_file=None):
        """Returns secret value from the AWS Secrets Manager or Parameter Store.

        :param secrets_store: the underlying secrets storage to be used, e.g. `secretsmanager` or `parameterstore`
        :type secrets_store: str
        :param secret_key: key for the secret value, e.g. `api-key`, `cert-key`, `trust-store`, `key-store`, defaults to None
        :type secret_key: str, optional
        :param api_key_file: path to a file containing an api key, defaults to None
        :type api_key_file: str optional
        :raises hsfs.client.exceptions.ExternalClientError: `api_key_file` needs to be set for local mode
        :raises hsfs.client.exceptions.UnknownSecretStorageError: Provided secrets storage not supported
        :return: secret
        :rtype: str
        """
        if secrets_store == self.SECRETS_MANAGER:
            return self._query_secrets_manager(secret_key)
        elif secrets_store == self.PARAMETER_STORE:
            return self._query_parameter_store(secret_key)
        elif secrets_store == self.LOCAL_STORE:
            if not api_key_file:
                raise exceptions.ExternalClientError(
                    "api_key_file needs to be set for local mode"
                )
            with open(api_key_file) as f:
                return f.readline().strip()
        else:
            raise exceptions.UnknownSecretStorageError(
                "Secrets storage " + secrets_store + " is not supported."
            )

    def _query_secrets_manager(self, secret_key):
        secret_name = "hopsworks/role/" + self._assumed_role()
        args = {"service_name": "secretsmanager"}
        region_name = self._get_region()
        if region_name:
            args["region_name"] = region_name
        client = boto3.client(**args)
        get_secret_value_response = client.get_secret_value(SecretId=secret_name)
        return json.loads(get_secret_value_response["SecretString"])[secret_key]

    def _assumed_role(self):
        client = boto3.client("sts")
        response = client.get_caller_identity()
        # arns for assumed roles in SageMaker follow the following schema
        # arn:aws:sts::123456789012:assumed-role/my-role-name/my-role-session-name
        local_identifier = response["Arn"].split(":")[-1].split("/")
        if len(local_identifier) != 3 or local_identifier[0] != "assumed-role":
            raise Exception(
                "Failed to extract assumed role from arn: " + response["Arn"]
            )
        return local_identifier[1]

    def _get_region(self):
        if self._region_name != self.DEFAULT_REGION:
            return self._region_name
        else:
            return None

    def _query_parameter_store(self, secret_key):
        args = {"service_name": "ssm"}
        region_name = self._get_region()
        if region_name:
            args["region_name"] = region_name
        client = boto3.client(**args)
        name = "/hopsworks/role/" + self._assumed_role() + "/type/" + secret_key
        return client.get_parameter(Name=name, WithDecryption=True)["Parameter"][
            "Value"
        ]

    def _get_project_info(self, project_name):
        """Makes a REST call to hopsworks to get all metadata of a project for the provided project.

        :param project_name: the name of the project
        :type project_name: str
        :return: JSON response with project info
        :rtype: dict
        """
        return self._send_request("GET", ["project", "getProjectInfo", project_name])

    def _write_b64_cert_to_bytes(self, b64_string, path):
        """Converts b64 encoded certificate to bytes file .

        :param b64_string:  b64 encoded string of certificate
        :type b64_string: str
        :param path: path where file is saved, including file name. e.g. /path/key-store.jks
        :type path: str
        """

        with open(path, "wb") as f:
            cert_b64 = base64.b64decode(b64_string)
            f.write(cert_b64)

    def _cleanup_file(self, file_path):
        """Removes local files with `file_path`."""
        try:
            os.remove(file_path)
        except OSError:
            pass

    def replace_public_host(self, url):
        """no need to replace as we are already in external client"""
        return url

    def _is_external(self):
        return True

    @property
    def host(self):
        return self._host
