# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.

"""
Credential classes for various cloud providers and services.
"""

import json
import os
import re
import sys
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Union

from packaging.requirements import InvalidRequirement, Requirement
from typeguard import typechecked


def validate_pypi_package(package: str) -> str:
    """
    Validate a PyPI package specifier using PEP 508 parsing.

    Args:
        package: Package specifier (e.g., "pandas>=2.0.0", "numpy[extra]")

    Returns:
        The validated package string.

    Raises:
        ValueError: If the package specifier is invalid.
    """
    package = package.strip()
    if not package:
        raise ValueError("Package specifier cannot be empty")
    try:
        Requirement(package)
        return package
    except InvalidRequirement as e:
        raise ValueError(f"Invalid PyPI package specifier '{package}': {e}") from e


def read_file(file_path: str) -> str:
    """Read file content as string."""
    with open(file_path, "r") as f:
        return f.read()


@typechecked
@dataclass
class AWSStaticCredentials:
    """
    Static Credentials for an AWS services.

    **Example**
    ```
    aws_credentials = ff.AWSStaticCredentials(
        access_key="<AWS_ACCESS_KEY>",
        secret_key="<AWS_SECRET_KEY>"
    )
    ```

    Args:
        access_key (str): AWS Access Key.
        secret_key (str): AWS Secret Key.
    """

    access_key: str = field(default="")
    secret_key: str = field(default="")
    _type: str = field(default="AWS_STATIC_CREDENTIALS")

    def __post_init__(
        self,
    ):
        if self.access_key == "":
            raise Exception("'AWSStaticCredentials' access_key cannot be empty")

        if self.secret_key == "":
            raise Exception("'AWSStaticCredentials' secret_key cannot be empty")

    @staticmethod
    def type() -> str:
        return "AWS_STATIC_CREDENTIALS"

    def config(self):
        return {
            "AccessKeyId": self.access_key,
            "SecretKey": self.secret_key,
            "Type": self._type,
        }


@typechecked
@dataclass
class AWSAssumeRoleCredentials:
    """
    Assume Role Credentials for an AWS services.

    If an IAM role for service account (IRSA) has been configured, the default credentials provider chain
        will be used to get the credentials stored on the pod. See the following link for more information:
        https://docs.aws.amazon.com/eks/latest/userguide/iam-roles-for-service-accounts.html

    **Example**
    ```
    aws_credentials = ff.AWSAssumeRoleCredentials()
    ```
    """

    _type: str = field(default="AWS_ASSUME_ROLE_CREDENTIALS")

    @staticmethod
    def type() -> str:
        return "AWS_ASSUME_ROLE_CREDENTIALS"

    def config(self):
        return {
            "Type": self._type,
        }


@typechecked
@dataclass
class GCPCredentials:
    def __init__(
        self, project_id: str, credentials_path: str = "", credential_json: dict = None
    ):
        """
        Credentials for an GCP.

        **Example**
        ```
        gcp_credentials = ff.GCPCredentials(
            project_id="<project_id>",
            credentials_path="<path_to_credentials>"
        )
        ```

        Args:
            project_id (str): The project id.
            credentials_path (str): The path to the credentials file.
        """
        if project_id == "":
            raise Exception("'GCPCredentials' project_id cannot be empty")

        self.project_id = project_id
        self.credentials_path = credentials_path

        if self.credentials_path == "" and credential_json is None:
            raise ValueError(
                "Either a path to credentials or json credentials must be provided"
            )
        elif self.credentials_path != "" and credential_json is not None:
            raise ValueError(
                "Only one of credentials_path or credentials_json can be provided"
            )
        elif self.credentials_path != "" and credential_json is None:
            if not os.path.isfile(credentials_path):
                raise Exception(
                    f"'GCPCredentials' credentials_path '{credentials_path}' file not found"
                )
            with open(credentials_path) as f:
                self.credentials = json.load(f)
        else:
            self.credentials = credential_json

    def type(self):
        return "GCPCredentials"

    def config(self):
        return {
            "ProjectId": self.project_id,
            "JSON": self.credentials,
        }

    def to_json(self):
        return self.credentials


@typechecked
@dataclass
class BasicCredentials:
    def __init__(
        self,
        username: str,
        password: str = "",
    ):
        """
        Credentials for an EMR cluster.

        **Example**
        ```
        creds = ff.BasicCredentials(
            username="<username>",
            password="<password>",
        )

        hdfs = ff.register_hdfs(
            name="hdfs",
            credentials=creds,
            ...
        )
        ```

        Args:
            username (str): the username of account.
            password (str): the password of account (optional).
        """
        self.username = username
        self.password = password

    def type(self):
        return "BasicCredential"

    def config(self):
        return {
            "Username": self.username,
            "Password": self.password,
        }


@typechecked
@dataclass
class KerberosCredentials:
    def __init__(
        self,
        username: str,
        password: str,
        krb5s_file: str,
    ):
        """
        Credentials for an EMR cluster.

        **Example**
        ```
        kerberos = ff.KerberosCredentials(
            username="<username>",
            password="<password>",
            krb5s_file="<path/to/krb5s_file>",
        )

        hdfs = ff.register_hdfs(
            name="hdfs",
            credentials=kerberos,
            ...
        )
        ```

        Args:
            username (str): the username of account.
            password (str): the password of account.
            krb5s_file (str): the path to krb5s file.
        """
        self.username = username
        self.password = password
        self.krb5s_conf = read_file(krb5s_file)

    def type(self):
        return "KerberosCredential"

    def config(self):
        return {
            "Username": self.username,
            "Password": self.password,
            "Krb5Conf": self.krb5s_conf,
        }


# Constants for Pyspark Versions
MAJOR_VERSION = "3"
MINOR_VERSIONS = ["9", "10", "11", "12"]


@typechecked
@dataclass
class DatabricksAwsAttributes:
    """
    AWS-specific settings for Databricks Jobs Compute clusters.

    Args:
        instance_profile_arn (str): IAM instance profile ARN attached to cluster nodes.
        availability (str): Availability mode such as ``ON_DEMAND`` or ``SPOT``.
        first_on_demand (int): Number of initial nodes placed on on-demand instances.
        spot_bid_price_percent (int): Spot bid price as a percentage of on-demand pricing.
        zone_id (str): Availability zone, for example ``us-west-2a``.
        ebs_volume_count (int): Number of EBS volumes attached per instance.
        ebs_volume_size (int): Size in GiB of each EBS volume.
        ebs_volume_type (str): EBS volume type, for example ``GENERAL_PURPOSE_SSD``.
        ebs_volume_iops (int): IOPS for gp3 volumes.
        ebs_volume_throughput (int): Throughput for gp3 volumes.
    """

    instance_profile_arn: str = ""
    availability: str = ""
    first_on_demand: int = 0
    spot_bid_price_percent: int = 0
    zone_id: str = ""
    ebs_volume_count: int = 0
    ebs_volume_size: int = 0
    ebs_volume_type: str = ""
    ebs_volume_iops: int = 0
    ebs_volume_throughput: int = 0

    def config(self):
        return {
            "InstanceProfileArn": self.instance_profile_arn,
            "Availability": self.availability,
            "FirstOnDemand": self.first_on_demand,
            "SpotBidPricePercent": self.spot_bid_price_percent,
            "ZoneId": self.zone_id,
            "EbsVolumeCount": self.ebs_volume_count,
            "EbsVolumeSize": self.ebs_volume_size,
            "EbsVolumeType": self.ebs_volume_type,
            "EbsVolumeIops": self.ebs_volume_iops,
            "EbsVolumeThroughput": self.ebs_volume_throughput,
        }


@typechecked
@dataclass
class DatabricksClusterConfig:
    """
    Configuration for an ephemeral Databricks job cluster.

    This configuration is used to create a new cluster for each job run,
    which is more cost-effective than using an always-on cluster.

    **Example**
    ```
    cluster_config = ff.DatabricksClusterConfig(
        spark_version="14.3.x-scala2.12",
        node_type_id="i3.xlarge",
        num_workers=4,
        libraries=["pandas==2.0.0", "numpy>=1.21.0"],
        aws_attributes=ff.DatabricksAwsAttributes(
            instance_profile_arn="arn:aws:iam::123456789012:instance-profile/featureform-databricks",
        ),
    )
    ```

    Args:
        spark_version (str): Required. Databricks Runtime version (e.g., "14.3.x-scala2.12").
        node_type_id (str): Required. Instance type for worker nodes (e.g., "i3.xlarge").
        num_workers (int): Number of worker nodes. 0 creates a single-node cluster.
        driver_node_type_id (str): Instance type for driver node. Defaults to node_type_id.
        spark_conf (Dict[str, str]): Cluster-level Spark configuration.
        custom_tags (Dict[str, str]): Custom tags for the cluster.
        init_scripts (List[str]): Init script paths.
        libraries (List[str]): PyPI package dependencies (e.g., ["pandas==2.0.0"]).
        aws_attributes (DatabricksAwsAttributes): AWS-specific cluster settings such as instance profile and spot behavior.
    """

    spark_version: str
    node_type_id: str
    num_workers: int = 2
    driver_node_type_id: str = ""
    spark_conf: Optional[Dict[str, str]] = None
    custom_tags: Optional[Dict[str, str]] = None
    init_scripts: Optional[List[str]] = None
    libraries: Optional[List[str]] = None
    aws_attributes: Optional[DatabricksAwsAttributes] = None

    def __post_init__(self):
        if self.libraries:
            for lib in self.libraries:
                validate_pypi_package(lib)


@typechecked
@dataclass
class DatabricksCredentials:
    """
    Credentials for a Databricks cluster.

    Supports two modes:
    1. Existing cluster mode: Use cluster_id to run jobs on a pre-existing cluster.
    2. Job cluster mode: Use cluster_config to create ephemeral clusters per job.

    **Example (PAT token with existing cluster)**
    ```
    databricks = ff.DatabricksCredentials(
        host="<databricks_hostname>",
        token="<databricks_token>",
        cluster_id="<databricks_cluster>",
    )
    ```

    **Example (OAuth M2M with existing cluster)**
    ```
    databricks = ff.DatabricksCredentials(
        host="<databricks_hostname>",
        client_id="<service_principal_client_id>",
        client_secret="<oauth_secret>",
        cluster_id="<databricks_cluster>",
    )
    ```

    **Example (job cluster - more cost-effective)**
    ```
    databricks = ff.DatabricksCredentials(
        host="<databricks_hostname>",
        token="<databricks_token>",
        cluster_config=ff.DatabricksClusterConfig(
            spark_version="14.3.x-scala2.12",
            node_type_id="i3.xlarge",
            num_workers=4,
        ),
    )
    ```

    Args:
        username (str): Username for a Databricks cluster (legacy auth).
        password (str): Password for a Databricks cluster (legacy auth).
        host (str): The hostname of a Databricks cluster.
        token (str): The personal access token for a Databricks cluster.
        cluster_id (str): Optional. ID of an existing Databricks cluster (mutually exclusive with cluster_config).
        client_id (str): OAuth M2M client ID for service principal authentication.
        client_secret (str): OAuth M2M client secret for service principal authentication.
        cluster_config (DatabricksClusterConfig): Optional. Configuration for ephemeral job clusters (mutually exclusive with cluster_id).
    """

    username: str = ""
    password: str = ""
    host: str = ""
    token: str = ""
    cluster_id: Optional[str] = None
    client_id: str = ""
    client_secret: str = ""
    cluster_config: Optional[DatabricksClusterConfig] = None

    def __post_init__(self):
        # Auth method 1: host + token (PAT)
        pat_auth = self.host != "" and self.token != ""
        # Auth method 2: username + password (legacy)
        basic_auth = self.username != "" and self.password != ""
        # Auth method 3: host + client_id + client_secret (OAuth M2M)
        oauth_auth = (
            self.host != "" and self.client_id != "" and self.client_secret != ""
        )

        auth_methods_count = sum([pat_auth, basic_auth, oauth_auth])

        if auth_methods_count != 1:
            raise ValueError(
                "DatabricksCredentials requires exactly one authentication method: "
                "('host' + 'token'), ('username' + 'password'), or "
                "('host' + 'client_id' + 'client_secret')"
            )

        # Mutual exclusivity: cluster_id XOR cluster_config
        if self.cluster_id and self.cluster_config:
            raise ValueError(
                "Cannot specify both cluster_id and cluster_config. "
                "Use cluster_id for existing clusters or cluster_config for job clusters."
            )

        if not self.cluster_id and not self.cluster_config:
            raise ValueError(
                "Must specify either cluster_id (existing cluster) or cluster_config (job cluster)."
            )

        # Validate cluster_id if provided
        if self.cluster_id:
            if not self._validate_cluster_id():
                raise ValueError(
                    f"Invalid cluster_id: expected id in the format 'xxxx-xxxxxx-xxxxxxxx' "
                    f"but received '{self.cluster_id}'"
                )

        # Validate cluster_config if provided
        if self.cluster_config:
            self._validate_cluster_config()

        # Validate token if using host/token auth
        if pat_auth and not self._validate_token():
            raise ValueError(
                f"Invalid token: expected token in the format 'dapi' + 32 alphanumeric characters "
                f"(optionally ending with '-' and 1 alphanumeric character) but received '{self.token}'"
            )

    def _validate_cluster_id(self) -> bool:
        cluster_id_regex = r"^\w{4}-\w{6}-\w{8}$"
        return bool(re.match(cluster_id_regex, self.cluster_id))

    def _validate_token(self) -> bool:
        token_regex = r"^dapi[a-zA-Z0-9]{32}(-[a-zA-Z0-9])?$"
        return bool(re.match(token_regex, self.token))

    def _validate_cluster_config(self) -> None:
        """Validate cluster_config fields."""
        cfg = self.cluster_config

        if not cfg.spark_version:
            raise ValueError("spark_version is required in cluster_config")

        if not self._validate_spark_version_format(cfg.spark_version):
            raise ValueError(
                f"Invalid spark_version format: '{cfg.spark_version}'. "
                f"Expected format: 'X.Y.x-scalaA.B' (e.g., '14.3.x-scala2.12')"
            )

        if not cfg.node_type_id:
            raise ValueError("node_type_id is required in cluster_config")

        if cfg.num_workers < 0:
            raise ValueError(f"num_workers must be >= 0, got {cfg.num_workers}")

        # Validate library package names if provided
        if cfg.libraries:
            for package in cfg.libraries:
                if not self._validate_pypi_package_name(package):
                    raise ValueError(
                        f"Invalid PyPI package name: '{package}'. "
                        f"Package names must follow PEP 508 conventions "
                        f"(start/end with alphanumeric, may contain hyphens, underscores, dots)."
                    )

    def _validate_spark_version_format(self, version: str) -> bool:
        """Validate Spark version format: X.Y.x-scalaA.B"""
        version_regex = r"^\d+\.\d+\.x-scala\d+\.\d+$"
        return bool(re.match(version_regex, version))

    def _validate_pypi_package_name(self, package: str) -> bool:
        """
        Validate PyPI package name follows PEP 508 conventions.

        Valid formats:
        - Package name only: pandas, numpy
        - With version: pandas==2.0.0, numpy>=1.21.0
        - With extras: requests[security]
        - With hyphens/underscores: scikit-learn, google_cloud_storage

        Invalid formats:
        - Starts/ends with hyphen or underscore: -invalid, invalid_
        - Empty string
        """
        if not package:
            return False

        # Regex for PEP 508 package names with optional version specifiers
        # Package name: starts/ends with alphanumeric, may contain .-_ in middle
        # Optional extras: [extra1,extra2]
        # Optional version: ==1.0.0, >=1.0.0, etc.
        pattern = (
            r"^[a-zA-Z0-9]([a-zA-Z0-9._-]*[a-zA-Z0-9])?"
            r"(\[.*\])?"
            r"((==|>=|<=|~=|!=|>|<|@).*)?$"
        )
        return bool(re.match(pattern, package))

    def type(self):
        return "DATABRICKS"

    def config(self):
        config = {
            "Username": self.username,
            "Password": self.password,
            "Host": self.host,
            "Token": self.token,
            "Cluster": self.cluster_id or "",
            "ClientID": self.client_id,
            "ClientSecret": self.client_secret,
        }

        # Add cluster_config fields if using job cluster mode
        if self.cluster_config:
            config["SparkVersion"] = self.cluster_config.spark_version
            config["NodeTypeId"] = self.cluster_config.node_type_id
            config["NumWorkers"] = self.cluster_config.num_workers
            config["DriverNodeTypeId"] = self.cluster_config.driver_node_type_id or ""
            config["SparkConf"] = self.cluster_config.spark_conf or {}
            config["CustomTags"] = self.cluster_config.custom_tags or {}
            config["InitScripts"] = self.cluster_config.init_scripts or []
            config["Libraries"] = self.cluster_config.libraries or []
            config["AwsAttributes"] = (
                self.cluster_config.aws_attributes.config()
                if self.cluster_config.aws_attributes
                else {}
            )

        return config


@typechecked
@dataclass
class EMRCredentials:
    def __init__(
        self,
        emr_cluster_id: str,
        emr_cluster_region: str,
        credentials: Union[AWSStaticCredentials, AWSAssumeRoleCredentials],
    ):
        """
        Credentials for an EMR cluster.

        **Example**
        ```
        emr = ff.EMRCredentials(
            emr_cluster_id="<cluster_id>",
            emr_cluster_region="<cluster_region>",
            credentials="<AWS_Credentials>",
        )

        spark = ff.register_spark(
            name="spark",
            executor=emr,
            ...
        )
        ```

        Args:
            emr_cluster_id (str): ID of an existing EMR cluster.
            emr_cluster_region (str): Region of an existing EMR cluster.
            credentials (Union[AWSStaticCredentials, AWSAssumeRoleCredentials]): Credentials for an AWS account with access to the cluster
        """
        self.emr_cluster_id = emr_cluster_id
        self.emr_cluster_region = emr_cluster_region
        self.credentials = credentials

    def type(self):
        return "EMR"

    def config(self):
        return {
            "ClusterName": self.emr_cluster_id,
            "ClusterRegion": self.emr_cluster_region,
            "Credentials": self.credentials.config(),
        }


@typechecked
@dataclass
class EMRServerlessCredentials:
    def __init__(
        self,
        application_id: str,
        region: str,
        execution_role_arn: str,
        credentials: Union[AWSStaticCredentials, AWSAssumeRoleCredentials],
        spark_submit_parameters: str = "",
        libraries: Optional[List[str]] = None,
    ):
        """
        Credentials for an EMR Serverless application.

        Args:
            application_id (str): ID of an existing EMR Serverless application.
            region (str): AWS region where the application exists.
            execution_role_arn (str): IAM role ARN used to execute job runs.
            credentials (Union[AWSStaticCredentials, AWSAssumeRoleCredentials]): Credentials for an AWS account with access to the application
            spark_submit_parameters (str): Optional spark-submit parameters applied to every job run.
            libraries (Optional[List[str]]): Optional extra PyPI dependencies to package with the job.
        """
        if application_id == "":
            raise Exception("'EMRServerlessCredentials' application_id cannot be empty")
        if region == "":
            raise Exception("'EMRServerlessCredentials' region cannot be empty")
        if execution_role_arn == "":
            raise Exception(
                "'EMRServerlessCredentials' execution_role_arn cannot be empty"
            )

        self.application_id = application_id
        self.region = region
        self.execution_role_arn = execution_role_arn
        self.credentials = credentials
        self.spark_submit_parameters = spark_submit_parameters
        self.libraries = libraries or []

        for lib in self.libraries:
            validate_pypi_package(lib)

    def type(self):
        return "EMR_SERVERLESS"

    def config(self):
        return {
            "ApplicationId": self.application_id,
            "Region": self.region,
            "ExecutionRoleARN": self.execution_role_arn,
            "Credentials": self.credentials.config(),
            "SparkSubmitParameters": self.spark_submit_parameters,
            "Libraries": self.libraries,
        }


@typechecked
@dataclass
class SparkCredentials:
    def __init__(
        self,
        master: str,
        deploy_mode: str,
        python_version: str,
        core_site_path: str = "",
        yarn_site_path: str = "",
    ):
        """
        Credentials for a Generic Spark Cluster

        **Example**
        ```
        spark_credentials = ff.SparkCredentials(
            master="yarn",
            deploy_mode="cluster",
            python_version="3.7.12",
            core_site_path="core-site.xml",
            yarn_site_path="yarn-site.xml"
        )

        spark = ff.register_spark(
            name="spark",
            executor=spark_credentials,
            ...
        )
        ```

        Args:
            master (str): The hostname of the Spark cluster. (The same that would be passed to `spark-submit`).
            deploy_mode (str): The deploy mode of the Spark cluster. (The same that would be passed to `spark-submit`).
            python_version (str): The Python version running on the cluster. Supports 3.9-3.12
            core_site_path (str): The path to the core-site.xml file. (For Yarn clusters only)
            yarn_site_path (str): The path to the yarn-site.xml file. (For Yarn clusters only)
        """
        self.master = master.lower()
        self.deploy_mode = deploy_mode.lower()
        self.core_site_path = core_site_path
        self.yarn_site_path = yarn_site_path

        if self.deploy_mode != "cluster" and self.deploy_mode != "client":
            raise Exception(
                f"Spark does not support '{self.deploy_mode}' deploy mode. It only supports 'cluster' and 'client'."
            )

        self.python_version = self._verify_python_version(
            self.deploy_mode, python_version
        )

        self._verify_yarn_config()

    def _verify_python_version(self, deploy_mode, version):
        if deploy_mode == "cluster" and version == "":
            client_python_version = sys.version_info
            client_major = str(client_python_version.major)
            client_minor = str(client_python_version.minor)

            if client_major != MAJOR_VERSION:
                client_major = "3"
            if client_minor not in MINOR_VERSIONS:
                client_minor = "7"

            version = f"{client_major}.{client_minor}"

        if version.count(".") == 2:
            major, minor, _ = version.split(".")
        elif version.count(".") == 1:
            major, minor = version.split(".")
        else:
            raise Exception(
                "Please specify your Python version on the Spark cluster. Accepted formats: Major.Minor or Major.Minor.Patch; ex. '3.9' or '3.9.16"
            )

        if major != MAJOR_VERSION or minor not in MINOR_VERSIONS:
            raise Exception(
                f"The Python version {version} is not supported. Currently, supported versions are 3.9-3.12."
            )

        """
        The Python versions on the Docker image are 3.9.16, 3.10.10, and 3.11.2.
        This conditional statement sets the patch number based on the minor version.
        """
        if minor == "10":
            patch = "10"
        elif minor == "11":
            patch = "2"
        else:
            patch = "16"

        return f"{major}.{minor}.{patch}"

    def _verify_yarn_config(self):
        if self.master == "yarn" and (
            self.core_site_path == "" or self.yarn_site_path == ""
        ):
            raise Exception(
                "Yarn requires core-site.xml and yarn-site.xml files."
                "Please copy these files from your Spark instance to local, then provide the local path in "
                "core_site_path and yarn_site_path. "
            )

    def type(self):
        return "SPARK"

    def config(self):
        core_site = (
            "" if self.core_site_path == "" else open(self.core_site_path, "r").read()
        )
        yarn_site = (
            "" if self.yarn_site_path == "" else open(self.yarn_site_path, "r").read()
        )

        return {
            "Master": self.master,
            "DeployMode": self.deploy_mode,
            "PythonVersion": self.python_version,
            "CoreSite": core_site,
            "YarnSite": yarn_site,
        }


# Type alias for executor credentials
ExecutorCredentials = Union[
    DatabricksCredentials,
    EMRCredentials,
    EMRServerlessCredentials,
    SparkCredentials,
]
