#  Copyright (c) ZenML GmbH 2022. 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.
"""SQL Model Implementations for Secrets."""

import base64
import json
from typing import Any, Dict, Optional, Sequence, cast
from uuid import UUID

from sqlalchemy import TEXT, VARCHAR, Column, UniqueConstraint
from sqlalchemy.orm import joinedload
from sqlalchemy.sql.base import ExecutableOption
from sqlalchemy_utils.types.encrypted.encrypted_type import (
    AesGcmEngine,
    InvalidCiphertextError,
)
from sqlmodel import Field, Relationship

from zenml.constants import TEXT_FIELD_MAX_LENGTH
from zenml.models import (
    SecretRequest,
    SecretResponse,
    SecretResponseBody,
    SecretResponseMetadata,
    SecretResponseResources,
    SecretUpdate,
)
from zenml.utils.time_utils import utc_now
from zenml.zen_stores.schemas.base_schemas import BaseSchema, NamedSchema
from zenml.zen_stores.schemas.schema_utils import (
    build_foreign_key_field,
    build_index,
)
from zenml.zen_stores.schemas.user_schemas import UserSchema
from zenml.zen_stores.schemas.utils import jl_arg


class SecretDecodeError(Exception):
    """Raised when a secret cannot be decoded or decrypted."""


class SecretSchema(NamedSchema, table=True):
    """SQL Model for secrets.

    Attributes:
        name: The name of the secret.
        values: The values of the secret.
    """

    __tablename__ = "secret"
    __table_args__ = (
        UniqueConstraint(
            "name",
            "private",
            "user_id",
            name="unique_secret_name_private_scope_user",
        ),
    )

    private: bool

    internal: bool = Field(default=False)

    values: Optional[bytes] = Field(sa_column=Column(TEXT, nullable=True))

    user_id: UUID = build_foreign_key_field(
        source=__tablename__,
        target=UserSchema.__tablename__,
        source_column="user_id",
        target_column="id",
        ondelete="CASCADE",
        nullable=False,
    )
    user: "UserSchema" = Relationship(back_populates="secrets")

    @classmethod
    def get_query_options(
        cls,
        include_metadata: bool = False,
        include_resources: bool = False,
        **kwargs: Any,
    ) -> Sequence[ExecutableOption]:
        """Get the query options for the schema.

        Args:
            include_metadata: Whether metadata will be included when converting
                the schema to a model.
            include_resources: Whether resources will be included when
                converting the schema to a model.
            **kwargs: Keyword arguments to allow schema specific logic

        Returns:
            A list of query options.
        """
        options = []

        if include_resources:
            options.extend([joinedload(jl_arg(SecretSchema.user))])

        return options

    @classmethod
    def _dump_secret_values(
        cls, values: Dict[str, str], encryption_engine: Optional[AesGcmEngine]
    ) -> bytes:
        """Dump the secret values to a string.

        Args:
            values: The secret values to dump.
            encryption_engine: The encryption engine to use to encrypt the
                secret values. If None, the values will be base64 encoded.

        Raises:
            ValueError: If the secret values do not fit in the database field.

        Returns:
            The serialized encrypted secret values.
        """
        serialized_values = json.dumps(values)

        if encryption_engine is None:
            encrypted_values = base64.b64encode(
                serialized_values.encode("utf-8")
            )
        else:
            encrypted_values = encryption_engine.encrypt(serialized_values)

        if len(encrypted_values) > TEXT_FIELD_MAX_LENGTH:
            raise ValueError(
                "Database representation of secret values exceeds max "
                "length. Please use fewer values or consider using shorter "
                "secret keys and/or values."
            )

        return encrypted_values

    @classmethod
    def _load_secret_values(
        cls,
        encrypted_values: bytes,
        encryption_engine: Optional[AesGcmEngine] = None,
    ) -> Dict[str, str]:
        """Load the secret values from a base64 encoded byte string.

        Args:
            encrypted_values: The serialized encrypted secret values.
            encryption_engine: The encryption engine to use to decrypt the
                secret values. If None, the values will be base64 decoded.

        Returns:
            The loaded secret values.

        Raises:
            SecretDecodeError: If the secret values cannot be decoded or
                decrypted.
        """
        if encryption_engine is None:
            try:
                serialized_values = base64.b64decode(encrypted_values).decode()
            except ValueError as e:
                raise SecretDecodeError(
                    "Could not decode base64 encoded secret values: {str(e)}"
                ) from e
        else:
            try:
                serialized_values = encryption_engine.decrypt(encrypted_values)
            except (ValueError, InvalidCiphertextError) as e:
                raise SecretDecodeError(
                    "Could not decrypt secret values. Please check that the "
                    f"encryption key is correct: {str(e)}"
                ) from e

        try:
            return cast(
                Dict[str, str],
                json.loads(serialized_values),
            )
        except json.JSONDecodeError as e:
            raise SecretDecodeError(
                "Could not decode secret values. Please check that the "
                f"secret values are valid JSON: {str(e)}"
            ) from e

    @classmethod
    def from_request(
        cls,
        secret: SecretRequest,
        internal: bool = False,
    ) -> "SecretSchema":
        """Create a `SecretSchema` from a `SecretRequest`.

        Args:
            secret: The `SecretRequest` from which to create the schema.
            internal: Whether the secret is internal.

        Returns:
            The created `SecretSchema`.
        """
        assert secret.user is not None, "User must be set for secret creation."
        return cls(
            name=secret.name,
            private=secret.private,
            user_id=secret.user,
            # Don't store secret values implicitly in the secret. The
            # SQL secret store will call `store_secret_values` to store the
            # values separately if SQL is used as the secrets store.
            values=None,
            internal=internal,
        )

    def update(
        self,
        secret_update: SecretUpdate,
    ) -> "SecretSchema":
        """Update a `SecretSchema` from a `SecretUpdate`.

        Args:
            secret_update: The `SecretUpdate` from which to update the schema.

        Returns:
            The updated `SecretSchema`.
        """
        # Don't update the secret values implicitly in the secret. The
        # SQL secret store will call `set_secret_values` to update the
        # values separately if SQL is used as the secrets store.
        for field, value in secret_update.model_dump(
            exclude_unset=True, exclude={"user", "values"}
        ).items():
            setattr(self, field, value)

        self.updated = utc_now()
        return self

    def to_model(
        self,
        include_metadata: bool = False,
        include_resources: bool = False,
        **kwargs: Any,
    ) -> SecretResponse:
        """Converts a secret schema to a secret model.

        Args:
            include_metadata: Whether the metadata will be filled.
            include_resources: Whether the resources will be filled.
            **kwargs: Keyword arguments to allow schema specific logic


        Returns:
            The secret model.
        """
        metadata = None
        if include_metadata:
            metadata = SecretResponseMetadata()

        resources = None
        if include_resources:
            resources = SecretResponseResources(
                user=self.user.to_model() if self.user else None,
            )

        # Don't load the secret values implicitly in the secret. The
        # SQL secret store will call `get_secret_values` to load the
        # values separately if SQL is used as the secrets store.
        body = SecretResponseBody(
            user_id=self.user_id,
            created=self.created,
            updated=self.updated,
            private=self.private,
        )
        return SecretResponse(
            id=self.id,
            name=self.name,
            body=body,
            metadata=metadata,
            resources=resources,
        )

    def get_secret_values(
        self,
        encryption_engine: Optional[AesGcmEngine] = None,
    ) -> Dict[str, str]:
        """Get the secret values for this secret.

        This method is used by the SQL secrets store to load the secret values
        from the database.

        Args:
            encryption_engine: The encryption engine to use to decrypt the
                secret values. If None, the values will be base64 decoded.

        Returns:
            The secret values

        Raises:
            KeyError: if no secret values for the given ID are stored in the
                secrets store.
        """
        if not self.values:
            raise KeyError(
                f"Secret values for secret {self.id} have not been stored in "
                f"the SQL secrets store."
            )
        return self._load_secret_values(self.values, encryption_engine)

    def set_secret_values(
        self,
        secret_values: Dict[str, str],
        encryption_engine: Optional[AesGcmEngine] = None,
    ) -> None:
        """Create a `SecretSchema` from a `SecretRequest`.

        This method is used by the SQL secrets store to store the secret values
        in the database.

        Args:
            secret_values: The new secret values.
            encryption_engine: The encryption engine to use to encrypt the
                secret values. If None, the values will be base64 encoded.
        """
        self.values = self._dump_secret_values(
            secret_values, encryption_engine
        )


class SecretResourceSchema(BaseSchema, table=True):
    """SQL Model for secret resource relationship."""

    __tablename__ = "secret_resource"
    __table_args__ = (
        build_index(
            table_name=__tablename__,
            column_names=[
                "resource_id",
                "resource_type",
                "secret_id",
            ],
            unique=True,
        ),
    )

    secret_id: UUID = build_foreign_key_field(
        source=__tablename__,
        target=SecretSchema.__tablename__,
        source_column="secret_id",
        target_column="id",
        ondelete="CASCADE",
        nullable=False,
    )
    resource_id: UUID
    resource_type: str = Field(sa_column=Column(VARCHAR(255), nullable=False))
