#  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.
"""Endpoint definitions for API keys."""

from typing import Union
from uuid import UUID

from fastapi import APIRouter, Depends, Security

from zenml.constants import (
    API,
    API_KEY_ROTATE,
    API_KEYS,
    SERVICE_ACCOUNTS,
    VERSION_1,
)
from zenml.exceptions import IllegalOperationError
from zenml.logger import get_logger
from zenml.models import (
    APIKeyFilter,
    APIKeyRequest,
    APIKeyResponse,
    APIKeyRotateRequest,
    APIKeyUpdate,
    Page,
    ServiceAccountFilter,
    ServiceAccountRequest,
    ServiceAccountResponse,
    ServiceAccountUpdate,
)
from zenml.models.v2.misc.server_models import ServerDeploymentType
from zenml.zen_server.auth import AuthContext, authorize
from zenml.zen_server.exceptions import error_response
from zenml.zen_server.rbac.endpoint_utils import (
    verify_permissions_and_create_entity,
    verify_permissions_and_get_entity,
    verify_permissions_and_list_entities,
)
from zenml.zen_server.rbac.models import Action, ResourceType
from zenml.zen_server.rbac.utils import (
    delete_model_resource,
    verify_permission_for_model,
)
from zenml.zen_server.utils import (
    async_fastapi_endpoint_wrapper,
    make_dependable,
    zen_store,
)

logger = get_logger(__name__)

router = APIRouter(
    prefix=API + VERSION_1 + SERVICE_ACCOUNTS,
    tags=["service_accounts", "api_keys"],
    responses={401: error_response},
)


def _raise_deprecated_pro_service_accounts() -> None:
    """Logs a warning if Pro workspace level service accounts are used."""
    from zenml.config.server_config import ServerConfiguration

    config = ServerConfiguration.get_server_config()
    if config.deployment_type == ServerDeploymentType.CLOUD:
        logger.warning(
            "ZenML Pro workspace level service accounts and API keys are "
            "deprecated and will be removed in a future version. Please use "
            "ZenML Pro organization level service accounts and API keys "
            "instead (see https://docs.zenml.io/pro/access-management/service-accounts)."
        )


# ----------------
# Service Accounts
# ----------------


@router.post(
    "",
    responses={
        401: error_response,
        409: error_response,
        422: error_response,
    },
)
@async_fastapi_endpoint_wrapper
def create_service_account(
    service_account: ServiceAccountRequest,
    _: AuthContext = Security(authorize),
) -> ServiceAccountResponse:
    """Creates a service account.

    Args:
        service_account: Service account to create.

    Returns:
        The created service account.
    """
    _raise_deprecated_pro_service_accounts()
    return verify_permissions_and_create_entity(
        request_model=service_account,
        create_method=zen_store().create_service_account,
    )


@router.get(
    "/{service_account_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def get_service_account(
    service_account_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
    hydrate: bool = True,
) -> ServiceAccountResponse:
    """Returns a specific service account.

    Args:
        service_account_name_or_id: Name or ID of the service account.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The service account matching the given name or ID.
    """
    return verify_permissions_and_get_entity(
        id=service_account_name_or_id,
        get_method=zen_store().get_service_account,
        hydrate=hydrate,
    )


@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def list_service_accounts(
    filter_model: ServiceAccountFilter = Depends(
        make_dependable(ServiceAccountFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ServiceAccountResponse]:
    """Returns a list of service accounts.

    Args:
        filter_model: Model that takes care of filtering, sorting and
            pagination.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        A list of service accounts matching the filter.
    """
    return verify_permissions_and_list_entities(
        filter_model=filter_model,
        resource_type=ResourceType.SERVICE_ACCOUNT,
        list_method=zen_store().list_service_accounts,
        hydrate=hydrate,
    )


@router.put(
    "/{service_account_name_or_id}",
    responses={
        401: error_response,
        404: error_response,
        422: error_response,
    },
)
@async_fastapi_endpoint_wrapper
def update_service_account(
    service_account_name_or_id: Union[str, UUID],
    service_account_update: ServiceAccountUpdate,
    _: AuthContext = Security(authorize),
) -> ServiceAccountResponse:
    """Updates a specific service account.

    Args:
        service_account_name_or_id: Name or ID of the service account.
        service_account_update: the service account to use for the update.

    Returns:
        The updated service account.

    Raises:
        IllegalOperationError: If the service account was created via external
            authentication.
    """
    _raise_deprecated_pro_service_accounts()
    service_account = zen_store().get_service_account(
        service_account_name_or_id, hydrate=True
    )
    if service_account.external_user_id is not None:
        raise IllegalOperationError(
            "Service accounts created via external authentication cannot "
            "be updated."
        )

    verify_permission_for_model(service_account, action=Action.UPDATE)

    return zen_store().update_service_account(
        service_account_name_or_id, service_account_update
    )


@router.delete(
    "/{service_account_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def delete_service_account(
    service_account_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
) -> None:
    """Delete a specific service account.

    Args:
        service_account_name_or_id: Name or ID of the service account.

    Raises:
        IllegalOperationError: If the service account was created via external
            authentication.
    """
    service_account = zen_store().get_service_account(
        service_account_name_or_id, hydrate=True
    )
    if service_account.external_user_id is not None:
        raise IllegalOperationError(
            "Service accounts created via external authentication cannot "
            "be deleted."
        )

    verify_permission_for_model(service_account, action=Action.DELETE)

    zen_store().delete_service_account(service_account_name_or_id)

    delete_model_resource(service_account)


# --------
# API Keys
# --------


@router.post(
    "/{service_account_id}" + API_KEYS,
    responses={401: error_response, 409: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def create_api_key(
    service_account_id: UUID,
    api_key: APIKeyRequest,
    _: AuthContext = Security(authorize),
) -> APIKeyResponse:
    """Creates an API key for a service account.

    Args:
        service_account_id: ID of the service account for which to create the
            API key.
        api_key: API key to create.

    Returns:
        The created API key.

    Raises:
        IllegalOperationError: If the service account was created via external
            authentication.
    """

    def create_api_key_wrapper(
        api_key: APIKeyRequest,
    ) -> APIKeyResponse:
        return zen_store().create_api_key(
            service_account_id=service_account_id,
            api_key=api_key,
        )

    _raise_deprecated_pro_service_accounts()
    service_account = zen_store().get_service_account(service_account_id)

    if service_account.external_user_id is not None:
        raise IllegalOperationError(
            "Service accounts created via external authentication cannot "
            "have associated API keys."
        )

    return verify_permissions_and_create_entity(
        request_model=api_key,
        create_method=create_api_key_wrapper,
        surrogate_models=[service_account],
    )


@router.get(
    "/{service_account_id}" + API_KEYS + "/{api_key_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def get_api_key(
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> APIKeyResponse:
    """Returns the requested API key.

    Args:
        service_account_id: ID of the service account to which the API key
            belongs.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        api_key_name_or_id: Name or ID of the API key to return.

    Returns:
        The requested API key.
    """
    service_account = zen_store().get_service_account(service_account_id)
    verify_permission_for_model(service_account, action=Action.READ)
    api_key = zen_store().get_api_key(
        service_account_id=service_account_id,
        api_key_name_or_id=api_key_name_or_id,
        hydrate=hydrate,
    )
    return api_key


@router.get(
    "/{service_account_id}" + API_KEYS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def list_api_keys(
    service_account_id: UUID,
    filter_model: APIKeyFilter = Depends(make_dependable(APIKeyFilter)),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[APIKeyResponse]:
    """List API keys associated with a service account.

    Args:
        service_account_id: ID of the service account to which the API keys
            belong.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        filter_model: Filter model used for pagination, sorting,
            filtering

    Returns:
        All API keys matching the filter and associated with the supplied
        service account.
    """
    service_account = zen_store().get_service_account(service_account_id)
    verify_permission_for_model(service_account, action=Action.READ)
    return zen_store().list_api_keys(
        service_account_id=service_account_id,
        filter_model=filter_model,
        hydrate=hydrate,
    )


@router.put(
    "/{service_account_id}" + API_KEYS + "/{api_key_name_or_id}",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def update_api_key(
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    api_key_update: APIKeyUpdate,
    _: AuthContext = Security(authorize),
) -> APIKeyResponse:
    """Updates an API key for a service account.

    Args:
        service_account_id: ID of the service account to which the API key
            belongs.
        api_key_name_or_id: Name or ID of the API key to update.
        api_key_update: API key update.

    Returns:
        The updated API key.

    Raises:
        IllegalOperationError: If the service account was created via external
            authentication.
    """
    _raise_deprecated_pro_service_accounts()
    service_account = zen_store().get_service_account(service_account_id)

    if service_account.external_user_id is not None:
        raise IllegalOperationError(
            "Service accounts created via external authentication cannot "
            "have associated API keys."
        )

    verify_permission_for_model(service_account, action=Action.UPDATE)
    return zen_store().update_api_key(
        service_account_id=service_account_id,
        api_key_name_or_id=api_key_name_or_id,
        api_key_update=api_key_update,
    )


@router.put(
    "/{service_account_id}"
    + API_KEYS
    + "/{api_key_name_or_id}"
    + API_KEY_ROTATE,
    responses={401: error_response, 409: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def rotate_api_key(
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    rotate_request: APIKeyRotateRequest,
    _: AuthContext = Security(authorize),
) -> APIKeyResponse:
    """Rotate an API key.

    Args:
        service_account_id: ID of the service account to which the API key
            belongs.
        api_key_name_or_id: Name or ID of the API key to rotate.
        rotate_request: API key rotation request.

    Returns:
        The updated API key.

    Raises:
        IllegalOperationError: If the service account was created via external
            authentication.
    """
    _raise_deprecated_pro_service_accounts()
    service_account = zen_store().get_service_account(service_account_id)

    if service_account.external_user_id is not None:
        raise IllegalOperationError(
            "Service accounts created via external authentication cannot "
            "have associated API keys."
        )

    verify_permission_for_model(service_account, action=Action.UPDATE)
    return zen_store().rotate_api_key(
        service_account_id=service_account_id,
        api_key_name_or_id=api_key_name_or_id,
        rotate_request=rotate_request,
    )


@router.delete(
    "/{service_account_id}" + API_KEYS + "/{api_key_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def delete_api_key(
    service_account_id: UUID,
    api_key_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes an API key.

    Args:
        service_account_id: ID of the service account to which the API key
            belongs.
        api_key_name_or_id: Name or ID of the API key to delete.
    """
    service_account = zen_store().get_service_account(service_account_id)
    verify_permission_for_model(service_account, action=Action.UPDATE)
    zen_store().delete_api_key(
        service_account_id=service_account_id,
        api_key_name_or_id=api_key_name_or_id,
    )
