#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
#
#  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:
#
#       https://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.
"""Utility functions for handling environment variables."""

import contextlib
import os
import re
from abc import ABC, abstractmethod
from typing import (
    TYPE_CHECKING,
    Any,
    Dict,
    Iterator,
    List,
    Match,
    Optional,
    TypeVar,
    Union,
    cast,
)

from pydantic import BaseModel
from typing_extensions import Self

from zenml.client import Client
from zenml.logger import get_logger
from zenml.utils import string_utils

if TYPE_CHECKING:
    from zenml.config.pipeline_configurations import PipelineConfiguration
    from zenml.config.step_configurations import StepConfiguration
    from zenml.models import StackResponse
    from zenml.stack import Stack

logger = get_logger(__name__)

V = TypeVar("V", bound=Any)
ENV_VARIABLE_PLACEHOLDER_PATTERN = re.compile(pattern=r"\$\{([a-zA-Z0-9_]+)\}")

ENV_VAR_CHUNK_SUFFIX = "_CHUNK_"


def split_environment_variables(
    size_limit: int,
    env: Optional[Dict[str, str]] = None,
) -> None:
    """Split long environment variables into chunks.

    Splits the input environment variables with values that exceed the supplied
    maximum length into individual components. The input environment variables
    are modified in-place.

    Args:
        size_limit: Maximum length of an environment variable value.
        env: Input environment variables dictionary. If not supplied, the
            OS environment variables are used.

    Raises:
        RuntimeError: If an environment variable value is too large and requires
            more than 10 chunks.
    """
    if env is None:
        env = cast(Dict[str, str], os.environ)

    for key, value in env.copy().items():
        if len(value) <= size_limit:
            continue

        # We keep the number of chunks to a maximum of 10 to avoid generating
        # too many environment variables chunks and also to make the
        # reconstruction easier to implement
        if len(value) > size_limit * 10:
            raise RuntimeError(
                f"Environment variable {key} exceeds the maximum length of "
                f"{size_limit * 10} characters."
            )

        env.pop(key)

        # Split the environment variable into chunks
        chunks = [
            value[i : i + size_limit] for i in range(0, len(value), size_limit)
        ]
        for i, chunk in enumerate(chunks):
            env[f"{key}{ENV_VAR_CHUNK_SUFFIX}{i}"] = chunk


def reconstruct_environment_variables(
    env: Optional[Dict[str, str]] = None,
) -> None:
    """Reconstruct environment variables that were split into chunks.

    Reconstructs the environment variables with values that were split into
    individual chunks because they were too large. The input environment
    variables are modified in-place.

    Args:
        env: Input environment variables dictionary. If not supplied, the OS
            environment variables are used.
    """
    if env is None:
        env = cast(Dict[str, str], os.environ)

    chunks: Dict[str, List[str]] = {}
    for key in env.keys():
        if not key[:-1].endswith(ENV_VAR_CHUNK_SUFFIX):
            continue

        # Collect all chunks of the same environment variable
        original_key = key[: -(len(ENV_VAR_CHUNK_SUFFIX) + 1)]
        chunks.setdefault(original_key, [])
        chunks[original_key].append(key)

    # Reconstruct the environment variables from their chunks
    for key, chunk_keys in chunks.items():
        chunk_keys.sort()
        value = "".join([env[key] for key in chunk_keys])
        env[key] = value

        # Remove the chunk environment variables
        for key in chunk_keys:
            env.pop(key)


def substitute_env_variable_placeholders(
    value: V, raise_when_missing: bool = True
) -> V:
    """Substitute environment variable placeholders in an object.

    Args:
        value: The object in which to substitute the placeholders.
        raise_when_missing: If True, an exception will be raised when an
            environment variable is missing. Otherwise, a warning will be logged
            instead.

    Returns:
        The object with placeholders substituted.
    """

    def _replace_with_env_variable_value(match: Match[str]) -> str:
        key = match.group(1)
        if key in os.environ:
            return os.environ[key]
        else:
            if raise_when_missing:
                raise KeyError(
                    "Unable to substitute environment variable placeholder "
                    f"'{key}' because the environment variable is not set."
                )
            else:
                logger.warning(
                    "Unable to substitute environment variable placeholder %s "
                    "because the environment variable is not set, using an "
                    "empty string instead.",
                    key,
                )
                return ""

    def _substitution_func(v: str) -> str:
        return ENV_VARIABLE_PLACEHOLDER_PATTERN.sub(
            _replace_with_env_variable_value, v
        )

    return string_utils.substitute_string(
        value=value, substitution_func=_substitution_func
    )


@contextlib.contextmanager
def temporary_environment(environment: Dict[str, str]) -> Iterator[None]:
    """Temporarily set environment variables.

    Args:
        environment: The environment variables to set.

    Yields:
        Nothing.
    """
    try:
        previous_env = {}
        for key, value in environment.items():
            previous_env[key] = os.environ.get(key, None)
            os.environ[key] = value
        yield
    finally:
        for key, previous_value in previous_env.items():
            updated_value = environment[key]
            current_value = os.environ.get(key)

            if current_value != updated_value:
                # The environment variable got updated while this context
                # manager was active -> We don't reset it back to the old
                # value
                continue
            else:
                if previous_value is None:
                    os.environ.pop(key, None)
                else:
                    os.environ[key] = previous_value


def get_runtime_environment(
    config: Union["PipelineConfiguration", "StepConfiguration"],
    stack: Union["Stack", "StackResponse"],
) -> Dict[str, str]:
    """Get the runtime environment variables.

    Args:
        config: The configuration.
        stack: The stack on which the pipeline/step will run.

    Returns:
        A dictionary of environment variables.
    """
    environment = {}
    for component in stack.components.values():
        if isinstance(component, list):
            for item in component:
                environment.update(item.environment)
        else:
            environment.update(component.environment)

    environment.update(stack.environment)
    environment.update(config.environment)

    for key, value in environment.items():
        environment[key] = str(value)

    return environment


def get_runtime_secret_environment(
    config: Union["PipelineConfiguration", "StepConfiguration"],
    stack: Union["Stack", "StackResponse"],
) -> Dict[str, str]:
    """Get the runtime secret environment variables.

    Args:
        config: The configuration.
        stack: The stack on which the pipeline/step will run.

    Returns:
        A dictionary of environment variables.
    """
    # We reverse the secrets to make sure the last defined secrets
    # override previous ones.
    # In the case of a step, the secrets will be the pipeline secrets followed
    # by the step secrets.
    secrets = list(reversed(config.secrets))
    secrets.extend(stack.secrets)

    for component in stack.components.values():
        if isinstance(component, list):
            for item in component:
                secrets.extend(item.secrets)
        else:
            secrets.extend(component.secrets)

    # Removes duplicates while preserving order, only the first occurrence of
    # each secret will be kept. We then reverse the list to make sure the
    # overrides are applied in the correct order.
    secrets = list(reversed(dict.fromkeys(secrets)))

    environment = {}
    for secret_name_or_id in secrets:
        try:
            secret = Client().get_secret(secret_name_or_id)
        except Exception as e:
            logger.warning(
                "Failed to get secret `%s` with error: %s. Skipping setting "
                "environment variables for this secret.",
                secret_name_or_id,
                e,
            )
            continue

        if not secret.secret_values:
            logger.warning(
                "Did not find any secret values for secret `%s`. This might be "
                "because you do not have permissions to read the secret "
                "values. Skipping setting environment variables for this "
                "secret.",
                secret_name_or_id,
            )
            continue

        for key, value in secret.secret_values.items():
            environment[key] = str(value)

    return environment


@contextlib.contextmanager
def temporary_runtime_environment(
    config: Union["PipelineConfiguration", "StepConfiguration"],
    stack: Union["Stack", "StackResponse"],
) -> Iterator[None]:
    """Temporarily set runtime environment variables.

    Args:
        config: The configuration.
        stack: The stack on which the pipeline/step will run.

    Yields:
        Nothing.
    """
    env = get_runtime_environment(config, stack)
    secret_env = get_runtime_secret_environment(config, stack)
    env.update(secret_env)

    with temporary_environment(env):
        yield


class ConfigBase(BaseModel, ABC):
    """Base class for configuration objects. Useful for config inheritance."""

    @staticmethod
    @abstractmethod
    def prefixes() -> list[str]:
        """Prefixes to use for env var name resolution.

        Returns:
            A list of prefixes to use for env var name resolution.
        """
        raise NotImplementedError

    @classmethod
    def get_env_vars(cls) -> dict[str, str]:
        """Method resolving the class variables from the environment.

        Notes:
        - Beware of variable name conflicts when composing configs via
        inheritance. For instance, if you have DB and Cache env vars with a
        different prefix (e.g. {ZENML_REDIS_}CONN_URL, {ZENML_DB_}CONN_URL) this
        may lead to a conflict in naming when combining the 2 individual config
        objects (We get 2 variables named conn_url).
        - To avoid generating ambiguous configs the function detects duplicates and fails eagerly.
        You can avoid such behaviors by using more generic prefixes and more precise
        variable names (e.g. prefix = ZENML_, variable = db_conn_url) or by controlling
        inheritance (try not to combine too many config classes, especially with common variable names).
        - Why not use pydantic_settings.BaseSettings? Because - while neat - it doesn't provide the
        ability to compose config objects with different prefixes via inheritance.

        Returns:
            A dictionary of config variables.

        Raises:
            RuntimeError: If config has conflicting variable names.
        """
        sorted_prefixes = sorted(
            set(cls.prefixes()), key=lambda p: (-len(p), p)
        )
        logger.info("Config resolving prefixes: %s", sorted_prefixes)

        env_vars: dict[str, str] = os.environ.copy()
        parameters: dict[str, str] = {}
        sources: dict[str, str] = {}  # variable_name -> env key

        for prefix in sorted_prefixes:
            found_keys: set[str] = set()

            for key, value in env_vars.items():
                if not key.startswith(prefix):
                    continue

                found_keys.add(key)
                variable_name = key.removeprefix(prefix).lower()

                if variable_name in parameters:
                    raise RuntimeError(
                        f"Environment variable extraction failed for {cls.__name__}. "
                        f"Variable '{variable_name}' is not uniquely defined; it maps to both "
                        f"{sources[variable_name]!r} and {key!r}."
                    )

                parameters[variable_name] = value
                sources[variable_name] = key

            env_vars = {
                k: v for k, v in env_vars.items() if k not in found_keys
            }

        return parameters

    @classmethod
    def load_from_env(cls) -> Self:
        """Resolve the environment and instantiate a new instance of the class.

        Returns:
            An instance of the class.
        """
        env_vars = cls.get_env_vars()
        logger.info(
            "Config identified the following variables: %s",
            list(env_vars.keys()),
        )
        return cls(**env_vars)
