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

from typing import Optional, Union
from uuid import UUID

from fastapi import APIRouter, Depends, Security

from zenml.constants import (
    API,
    ARTIFACTS,
    MODEL_VERSION_ARTIFACTS,
    MODEL_VERSION_PIPELINE_RUNS,
    MODEL_VERSIONS,
    MODELS,
    RUNS,
    VERSION_1,
)
from zenml.models import (
    ModelVersionArtifactFilter,
    ModelVersionArtifactRequest,
    ModelVersionArtifactResponse,
    ModelVersionFilter,
    ModelVersionPipelineRunFilter,
    ModelVersionPipelineRunRequest,
    ModelVersionPipelineRunResponse,
    ModelVersionRequest,
    ModelVersionResponse,
    ModelVersionUpdate,
)
from zenml.models.v2.base.page import Page
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_delete_entity,
)
from zenml.zen_server.rbac.models import Action, ResourceType
from zenml.zen_server.rbac.utils import (
    dehydrate_page,
    dehydrate_response_model,
    get_allowed_resource_ids,
    verify_permission_for_model,
)
from zenml.zen_server.routers.models_endpoints import (
    router as model_router,
)
from zenml.zen_server.routers.projects_endpoints import workspace_router
from zenml.zen_server.utils import (
    async_fastapi_endpoint_wrapper,
    make_dependable,
    set_filter_project_scope,
    zen_store,
)

#################
# Model Versions
#################


router = APIRouter(
    prefix=API + VERSION_1 + MODEL_VERSIONS,
    tags=["model_versions"],
    responses={401: error_response, 403: error_response},
)


@router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
# TODO: the workspace scoped endpoint is only kept for dashboard compatibility
# and can be removed after the migration
@workspace_router.post(
    "/{project_name_or_id}" + MODELS + "/{model_id}" + MODEL_VERSIONS,
    responses={401: error_response, 409: error_response, 422: error_response},
    deprecated=True,
    tags=["model_versions"],
)
@async_fastapi_endpoint_wrapper
def create_model_version(
    model_version: ModelVersionRequest,
    model_id: Optional[UUID] = None,
    project_name_or_id: Optional[Union[str, UUID]] = None,
    _: AuthContext = Security(authorize),
) -> ModelVersionResponse:
    """Creates a model version.

    Args:
        model_version: Model version to create.
        model_id: Optional ID of the model.
        project_name_or_id: Optional name or ID of the project.

    Returns:
        The created model version.
    """
    if project_name_or_id:
        project = zen_store().get_project(project_name_or_id)
        model_version.project = project.id

    if model_id:
        model_version.model = model_id

    return verify_permissions_and_create_entity(
        request_model=model_version,
        create_method=zen_store().create_model_version,
    )


@model_router.get(
    "/{model_name_or_id}" + MODEL_VERSIONS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper(deduplicate=True)
def list_model_versions(
    model_version_filter_model: ModelVersionFilter = Depends(
        make_dependable(ModelVersionFilter)
    ),
    model_name_or_id: Optional[Union[str, UUID]] = None,
    hydrate: bool = False,
    auth_context: AuthContext = Security(authorize),
) -> Page[ModelVersionResponse]:
    """Get model versions according to query filters.

    Args:
        model_version_filter_model: Filter model used for pagination, sorting,
            filtering.
        model_name_or_id: Optional name or ID of the model.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.
        auth_context: The authentication context.

    Returns:
        The model versions according to query filters.
    """
    # A project scoped request must always be scoped to a specific
    # project. This is required for the RBAC check to work.
    set_filter_project_scope(model_version_filter_model)
    assert isinstance(model_version_filter_model.project, UUID)

    if model_name_or_id:
        model_version_filter_model.model = model_name_or_id

    allowed_model_ids = get_allowed_resource_ids(
        resource_type=ResourceType.MODEL,
        project_id=model_version_filter_model.project,
    )
    model_version_filter_model.configure_rbac(
        authenticated_user_id=auth_context.user.id,
        model_id=allowed_model_ids,
    )
    model_versions = zen_store().list_model_versions(
        model_version_filter_model=model_version_filter_model,
        hydrate=hydrate,
    )
    return dehydrate_page(model_versions)


@router.get(
    "/{model_version_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper(deduplicate=True)
def get_model_version(
    model_version_id: UUID,
    hydrate: bool = True,
    _: AuthContext = Security(authorize),
) -> ModelVersionResponse:
    """Get a model version by ID.

    Args:
        model_version_id: id of the model version to be retrieved.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The model version with the given name or ID.
    """
    model_version = zen_store().get_model_version(
        model_version_id=model_version_id,
        hydrate=hydrate,
    )
    verify_permission_for_model(model_version.model, action=Action.READ)
    return dehydrate_response_model(model_version)


@router.put(
    "/{model_version_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper(deduplicate=True)
def update_model_version(
    model_version_id: UUID,
    model_version_update_model: ModelVersionUpdate,
    _: AuthContext = Security(authorize),
) -> ModelVersionResponse:
    """Get all model versions by filter.

    Args:
        model_version_id: The ID of model version to be updated.
        model_version_update_model: The model version to be updated.

    Returns:
        An updated model version.
    """
    model_version = zen_store().get_model_version(model_version_id)

    if model_version_update_model.stage:
        # Make sure the user has permissions to promote the model
        verify_permission_for_model(model_version.model, action=Action.PROMOTE)

    verify_permission_for_model(model_version, action=Action.UPDATE)
    updated_model_version = zen_store().update_model_version(
        model_version_id=model_version_id,
        model_version_update_model=model_version_update_model,
    )

    return dehydrate_response_model(updated_model_version)


@router.delete(
    "/{model_version_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def delete_model_version(
    model_version_id: UUID,
    _: AuthContext = Security(authorize),
) -> None:
    """Delete a model by name or ID.

    Args:
        model_version_id: The name or ID of the model version to delete.
    """
    verify_permissions_and_delete_entity(
        id=model_version_id,
        get_method=zen_store().get_model_version,
        delete_method=zen_store().delete_model_version,
    )


##########################
# Model Version Artifacts
##########################

model_version_artifacts_router = APIRouter(
    prefix=API + VERSION_1 + MODEL_VERSION_ARTIFACTS,
    tags=["model_version_artifacts"],
    responses={401: error_response},
)


@model_version_artifacts_router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def create_model_version_artifact_link(
    model_version_artifact_link: ModelVersionArtifactRequest,
    _: AuthContext = Security(authorize),
) -> ModelVersionArtifactResponse:
    """Create a new model version to artifact link.

    Args:
        model_version_artifact_link: The model version to artifact link to create.

    Returns:
        The created model version to artifact link.
    """
    model_version = zen_store().get_model_version(
        model_version_artifact_link.model_version
    )

    return verify_permissions_and_create_entity(
        request_model=model_version_artifact_link,
        create_method=zen_store().create_model_version_artifact_link,
        # Check for UPDATE permissions on the model version instead of the
        # model version artifact link
        surrogate_models=[model_version],
    )


@model_version_artifacts_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def list_model_version_artifact_links(
    model_version_artifact_link_filter_model: ModelVersionArtifactFilter = Depends(
        make_dependable(ModelVersionArtifactFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ModelVersionArtifactResponse]:
    """Get model version to artifact links according to query filters.

    Args:
        model_version_artifact_link_filter_model: Filter model used for
            pagination, sorting, filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The model version to artifact links according to query filters.
    """
    return zen_store().list_model_version_artifact_links(
        model_version_artifact_link_filter_model=model_version_artifact_link_filter_model,
        hydrate=hydrate,
    )


@router.delete(
    "/{model_version_id}"
    + ARTIFACTS
    + "/{model_version_artifact_link_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def delete_model_version_artifact_link(
    model_version_id: UUID,
    model_version_artifact_link_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a model version to artifact link.

    Args:
        model_version_id: ID of the model version containing the link.
        model_version_artifact_link_name_or_id: name or ID of the model
            version to artifact link to be deleted.
    """
    model_version = zen_store().get_model_version(model_version_id)
    verify_permission_for_model(model_version, action=Action.UPDATE)

    zen_store().delete_model_version_artifact_link(
        model_version_id,
        model_version_artifact_link_name_or_id,
    )


@router.delete(
    "/{model_version_id}" + ARTIFACTS,
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def delete_all_model_version_artifact_links(
    model_version_id: UUID,
    only_links: bool = True,
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes all model version to artifact links.

    Args:
        model_version_id: ID of the model version containing links.
        only_links: Whether to only delete the link to the artifact.
    """
    model_version = zen_store().get_model_version(model_version_id)
    verify_permission_for_model(model_version, action=Action.UPDATE)

    zen_store().delete_all_model_version_artifact_links(
        model_version_id, only_links
    )


##############################
# Model Version Pipeline Runs
##############################

model_version_pipeline_runs_router = APIRouter(
    prefix=API + VERSION_1 + MODEL_VERSION_PIPELINE_RUNS,
    tags=["model_version_pipeline_runs"],
    responses={401: error_response},
)


@model_version_pipeline_runs_router.post(
    "",
    responses={401: error_response, 409: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def create_model_version_pipeline_run_link(
    model_version_pipeline_run_link: ModelVersionPipelineRunRequest,
    _: AuthContext = Security(authorize),
) -> ModelVersionPipelineRunResponse:
    """Create a new model version to pipeline run link.

    Args:
        model_version_pipeline_run_link: The model version to pipeline run link to create.

    Returns:
        - If Model Version to Pipeline Run Link already exists - returns the existing link.
        - Otherwise, returns the newly created model version to pipeline run link.
    """
    model_version = zen_store().get_model_version(
        model_version_pipeline_run_link.model_version, hydrate=False
    )

    return verify_permissions_and_create_entity(
        request_model=model_version_pipeline_run_link,
        create_method=zen_store().create_model_version_pipeline_run_link,
        # Check for UPDATE permissions on the model version instead of the
        # model version pipeline run link
        surrogate_models=[model_version],
    )


@model_version_pipeline_runs_router.get(
    "",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def list_model_version_pipeline_run_links(
    model_version_pipeline_run_link_filter_model: ModelVersionPipelineRunFilter = Depends(
        make_dependable(ModelVersionPipelineRunFilter)
    ),
    hydrate: bool = False,
    _: AuthContext = Security(authorize),
) -> Page[ModelVersionPipelineRunResponse]:
    """Get model version to pipeline run links according to query filters.

    Args:
        model_version_pipeline_run_link_filter_model: Filter model used for
            pagination, sorting, and filtering.
        hydrate: Flag deciding whether to hydrate the output model(s)
            by including metadata fields in the response.

    Returns:
        The model version to pipeline run links according to query filters.
    """
    return zen_store().list_model_version_pipeline_run_links(
        model_version_pipeline_run_link_filter_model=model_version_pipeline_run_link_filter_model,
        hydrate=hydrate,
    )


@router.delete(
    "/{model_version_id}"
    + RUNS
    + "/{model_version_pipeline_run_link_name_or_id}",
    responses={401: error_response, 404: error_response, 422: error_response},
)
@async_fastapi_endpoint_wrapper
def delete_model_version_pipeline_run_link(
    model_version_id: UUID,
    model_version_pipeline_run_link_name_or_id: Union[str, UUID],
    _: AuthContext = Security(authorize),
) -> None:
    """Deletes a model version link.

    Args:
        model_version_id: name or ID of the model version containing the link.
        model_version_pipeline_run_link_name_or_id: name or ID of the model
            version link to be deleted.
    """
    model_version = zen_store().get_model_version(model_version_id)
    verify_permission_for_model(model_version, action=Action.UPDATE)

    zen_store().delete_model_version_pipeline_run_link(
        model_version_id=model_version_id,
        model_version_pipeline_run_link_name_or_id=model_version_pipeline_run_link_name_or_id,
    )
