#  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.
"""CLI functionality to interact with deployments."""

import json
from typing import Any, List, Optional
from uuid import UUID

import click

from zenml.cli import utils as cli_utils
from zenml.cli.cli import TagGroup, cli
from zenml.cli.utils import OutputFormat, fetch_snapshot, list_options
from zenml.client import Client
from zenml.console import console
from zenml.deployers.exceptions import DeploymentInvalidParametersError
from zenml.enums import CliCategories
from zenml.logger import get_logger
from zenml.models import (
    DeploymentFilter,
)
from zenml.pipelines.pipeline_definition import Pipeline
from zenml.utils import source_utils
from zenml.utils.dashboard_utils import get_deployment_url

logger = get_logger(__name__)


def _import_pipeline(source: str) -> Pipeline:
    """Import a pipeline.

    Args:
        source: The pipeline source.

    Returns:
        The pipeline.
    """
    try:
        pipeline_instance = source_utils.load(source)
    except ModuleNotFoundError as e:
        source_root = source_utils.get_source_root()
        cli_utils.error(
            f"Unable to import module `{e.name}`. Make sure the source path is "
            f"relative to your source root `{source_root}`."
        )
    except AttributeError as e:
        cli_utils.error("Unable to load attribute from module: " + str(e))

    if not isinstance(pipeline_instance, Pipeline):
        cli_utils.error(
            f"The given source path `{source}` does not resolve to a pipeline "
            "object."
        )

    return pipeline_instance


@cli.group(cls=TagGroup, tag=CliCategories.MANAGEMENT_TOOLS)
def deployment() -> None:
    """Interact with deployments."""


@deployment.command("list", help="List all registered deployments.")
@list_options(
    DeploymentFilter,
    default_columns=[
        "id",
        "name",
        "status",
        "url",
        "pipeline",
        "stack",
    ],
)
def list_deployments(
    columns: str, output_format: OutputFormat, **kwargs: Any
) -> None:
    """List all registered deployments for the filter.

    Args:
        columns: Columns to display in output.
        output_format: Format for output (table/json/yaml/csv/tsv).
        **kwargs: Keyword arguments to filter deployments.
    """
    client = Client()
    try:
        with console.status("Listing deployments...\n"):
            deployments = client.list_deployments(**kwargs)
    except KeyError as err:
        cli_utils.exception(err)
        return

    cli_utils.print_page(
        deployments,
        columns,
        output_format,
        cli_utils.generate_deployment_row,
        empty_message="No deployments found for this filter.",
    )


@deployment.command("describe", help="Describe a deployment.")
@click.argument("deployment_name_or_id", type=str, required=True)
@click.option(
    "--show-secret",
    "-s",
    is_flag=True,
    help="Show the secret key.",
)
@click.option(
    "--show-metadata",
    "-m",
    is_flag=True,
    help="Show the metadata.",
)
@click.option(
    "--show-schema",
    "-sc",
    is_flag=True,
    help="Show the schema.",
)
@click.option(
    "--no-truncate",
    "-nt",
    is_flag=True,
    help="Don't truncate the metadata.",
)
def describe_deployment(
    deployment_name_or_id: str,
    show_secret: bool = False,
    show_metadata: bool = False,
    no_truncate: bool = False,
    show_schema: bool = False,
) -> None:
    """Describe a deployment.

    Args:
        deployment_name_or_id: The name or ID of the deployment to describe.
        show_secret: If True, show the secret key.
        show_metadata: If True, show the metadata.
        show_schema: If True, show the schema.
        no_truncate: If True, don't truncate the metadata.
    """
    try:
        deployment = Client().get_deployment(
            name_id_or_prefix=deployment_name_or_id,
        )
    except KeyError as e:
        cli_utils.exception(e)
    else:
        cli_utils.pretty_print_deployment(
            deployment,
            show_secret=show_secret,
            show_metadata=show_metadata,
            show_schema=show_schema,
            no_truncate=no_truncate,
        )


@deployment.command("provision", help="Provision a deployment.")
@click.argument("deployment_name_or_id", type=str, required=True)
@click.option(
    "--snapshot",
    "-s",
    "snapshot_name_or_id",
    type=str,
    required=False,
    help="ID or name of the snapshot to use. If not provided, the current "
    "snapshot configured for the deployment will be used.",
)
@click.option(
    "--pipeline",
    "-p",
    "pipeline_name_or_id",
    type=str,
    required=False,
    help="The name or ID of the pipeline to which the snapshot belongs.",
)
@click.option(
    "--overtake",
    "-o",
    "overtake",
    is_flag=True,
    default=False,
    required=False,
    help="Provision the deployment with the given name even if it is "
    "owned by a different user.",
)
@click.option(
    "--timeout",
    "-t",
    "timeout",
    type=int,
    required=False,
    default=None,
    help="Maximum time in seconds to wait for the deployment to be "
    "provisioned.",
)
def provision_deployment(
    deployment_name_or_id: str,
    snapshot_name_or_id: Optional[str] = None,
    pipeline_name_or_id: Optional[str] = None,
    overtake: bool = False,
    timeout: Optional[int] = None,
) -> None:
    """Provision a deployment.

    Args:
        deployment_name_or_id: The name or ID of the deployment to deploy.
        snapshot_name_or_id: The ID or name of the pipeline snapshot to use.
        pipeline_name_or_id: The name or ID of the pipeline to which the
            snapshot belongs.
        overtake: If True, provision the deployment with the given name
            even if it is owned by a different user.
        timeout: The maximum time in seconds to wait for the deployment
            to be provisioned.
    """
    snapshot_id: Optional[UUID] = None
    if snapshot_name_or_id:
        snapshot = fetch_snapshot(snapshot_name_or_id, pipeline_name_or_id)
        snapshot_id = snapshot.id

    client = Client()
    try:
        deployment = client.get_deployment(deployment_name_or_id)
    except KeyError:
        pass
    else:
        if (
            deployment.user
            and deployment.user.id != client.active_user.id
            and not overtake
        ):
            confirmation = cli_utils.confirmation(
                f"Deployment with name '{deployment.name}' is owned by a "
                f"different user '{deployment.user.name}'.\nDo you want to "
                "continue and provision it "
                "(hint: use the --overtake flag to skip this check)?"
            )
            if not confirmation:
                cli_utils.declare("Deployment provisioning canceled.")
                return

    with console.status(
        f"Provisioning deployment '{deployment_name_or_id}'...\n"
    ):
        try:
            deployment = Client().provision_deployment(
                name_id_or_prefix=deployment_name_or_id,
                snapshot_id=snapshot_id,
                timeout=timeout,
            )
        except KeyError as e:
            cli_utils.exception(e)
        else:
            cli_utils.declare(
                f"Provisioned deployment '{deployment_name_or_id}'."
            )
            cli_utils.pretty_print_deployment(deployment, show_secret=True)

            dashboard_url = get_deployment_url(deployment)
            if dashboard_url:
                cli_utils.declare(
                    f"\nView in ZenML Cloud: [link]{dashboard_url}[/link]",
                    no_wrap=True,
                )


@deployment.command("deprovision", help="Deprovision a deployment.")
@click.argument("deployment_name_or_id", type=str, required=False)
@click.option(
    "--all",
    "-a",
    is_flag=True,
    default=False,
    help="Deprovision all deployments.",
)
@click.option(
    "--mine",
    "-m",
    is_flag=True,
    default=False,
    help="Deprovision only deployments owned by the current user.",
)
@click.option(
    "--ignore-errors",
    "-i",
    is_flag=True,
    default=False,
    help="Ignore errors when deprovisioning multiple deployments.",
)
@click.option(
    "--yes",
    "-y",
    is_flag=True,
    default=False,
    help="Don't ask for confirmation.",
)
@click.option(
    "--max-count",
    "-c",
    "max_count",
    type=int,
    required=False,
    default=10,
    help="Maximum number of deployments to deprovision in one go.",
)
@click.option(
    "--timeout",
    "-t",
    "timeout",
    type=int,
    required=False,
    default=None,
    help="Maximum time in seconds to wait for the deployment to be "
    "deprovisioned.",
)
def deprovision_deployment(
    deployment_name_or_id: Optional[str] = None,
    all: bool = False,
    mine: bool = False,
    yes: bool = False,
    ignore_errors: bool = False,
    max_count: int = 10,
    timeout: Optional[int] = None,
) -> None:
    """Deprovision a deployment.

    Args:
        deployment_name_or_id: The name or ID of the deployment to deprovision.
        all: If set, deprovision all deployments.
        mine: If set, deprovision only deployments owned by the current user.
        yes: If set, don't ask for confirmation.
        ignore_errors: If set, ignore errors when deprovisioning multiple
            deployments.
        max_count: The maximum number of deployments to deprovision in one go.
        timeout: The maximum time in seconds to wait for the deployment
            to be deprovisioned.
    """
    client = Client()

    if all:
        deployments = client.list_deployments(size=max_count).items
    elif mine:
        deployments = client.list_deployments(
            user=client.active_user.id,
            size=max_count,
        ).items
    elif deployment_name_or_id:
        deployments = [
            client.get_deployment(name_id_or_prefix=deployment_name_or_id)
        ]
    else:
        cli_utils.error(
            "Either a deployment name or ID must be provided or --all or "
            "--mine must be set."
        )

    if len(deployments) == 0:
        cli_utils.error("No deployments found to deprovision.")

    if len(deployments) == 1:
        cli_utils.declare(
            f"The following deployment will be deprovisioned: "
            f"{deployments[0].name}"
        )
    else:
        deployment_names = [deployment.name for deployment in deployments]
        cli_utils.declare(
            f"The following deployments will ALL be deprovisioned: "
            f"{', '.join(deployment_names)}"
        )

    other_deployments = [
        deployment
        for deployment in deployments
        if deployment.user and deployment.user.id != client.active_user.id
    ]

    if other_deployments:
        deployment_names = [
            deployment.name for deployment in other_deployments
        ]
        cli_utils.warning(
            f"The following deployments are not owned by the current "
            f"user: {', '.join(deployment_names)} !"
        )

    if not yes:
        confirmation = cli_utils.confirmation(
            "Are you sure you want to continue?"
        )
        if not confirmation:
            cli_utils.declare("Deployment deprovision canceled.")
            return

    for deployment in deployments:
        with console.status(
            f"Deprovisioning deployment '{deployment.name}'...\n"
        ):
            try:
                client.deprovision_deployment(
                    name_id_or_prefix=deployment.id,
                    timeout=timeout,
                )
                cli_utils.declare(
                    f"Deprovisioned deployment '{deployment.name}'."
                )
                cli_utils.declare(
                    "Hint: to permanently delete the deployment, run `zenml "
                    f"deployment delete {deployment.name}`."
                )
            except KeyError as e:
                error_message = (
                    f"Failed to deprovision deployment '{deployment.name}': "
                    f"{str(e)}"
                )
                if all or mine and ignore_errors:
                    cli_utils.warning(error_message)
                else:
                    cli_utils.error(error_message)


@deployment.command("delete", help="Delete a deployment.")
@click.argument("deployment_name_or_id", type=str, required=False)
@click.option(
    "--all",
    "-a",
    is_flag=True,
    default=False,
    help="Deprovision all deployments.",
)
@click.option(
    "--mine",
    "-m",
    is_flag=True,
    default=False,
    help="Deprovision only deployments owned by the current user.",
)
@click.option(
    "--ignore-errors",
    "-i",
    is_flag=True,
    default=False,
    help="Ignore errors when deprovisioning multiple deployments.",
)
@click.option(
    "--yes",
    "-y",
    is_flag=True,
    default=False,
    help="Don't ask for confirmation.",
)
@click.option(
    "--timeout",
    "-t",
    "timeout",
    type=int,
    required=False,
    default=None,
    help="Maximum time in seconds to wait for the deployment to be "
    "deprovisioned.",
)
@click.option(
    "--max-count",
    "-c",
    "max_count",
    type=int,
    required=False,
    default=20,
    help="Maximum number of deployments to delete in one go.",
)
@click.option(
    "--force",
    "-f",
    is_flag=True,
    default=False,
    help="Force the deletion of the deployment if it cannot be deprovisioned.",
)
def delete_deployment(
    deployment_name_or_id: Optional[str] = None,
    all: bool = False,
    mine: bool = False,
    ignore_errors: bool = False,
    yes: bool = False,
    timeout: Optional[int] = None,
    max_count: int = 20,
    force: bool = False,
) -> None:
    """Deprovision and delete a deployment.

    Args:
        deployment_name_or_id: The name or ID of the deployment to delete.
        all: If set, delete all deployments.
        mine: If set, delete only deployments owned by the current user.
        ignore_errors: If set, ignore errors when deleting multiple deployments.
        yes: If set, don't ask for confirmation.
        timeout: The maximum time in seconds to wait for the deployment
            to be deprovisioned.
        max_count: The maximum number of deployments to delete in one go.
        force: If set, force the deletion of the deployment if it cannot
            be deprovisioned.
    """
    client = Client()

    if all:
        deployments = client.list_deployments(size=max_count).items
    elif mine:
        deployments = client.list_deployments(
            user=client.active_user.id,
            size=max_count,
        ).items
    elif deployment_name_or_id:
        deployments = [
            client.get_deployment(name_id_or_prefix=deployment_name_or_id)
        ]
    else:
        cli_utils.error(
            "Either a deployment name or ID must be provided or --all or "
            "--mine must be set."
        )

    if len(deployments) == 0:
        cli_utils.error("No deployments found to delete.")

    if len(deployments) == 1:
        cli_utils.declare(
            f"The following deployment will be deleted: {deployments[0].name}"
        )
    else:
        deployment_names = [deployment.name for deployment in deployments]
        cli_utils.declare(
            f"The following deployments will ALL be deleted: "
            f"{', '.join(deployment_names)}"
        )

    other_deployments = [
        deployment
        for deployment in deployments
        if deployment.user and deployment.user.id != client.active_user.id
    ]

    if other_deployments:
        deployment_names = [
            deployment.name for deployment in other_deployments
        ]
        cli_utils.warning(
            f"The following deployments are not owned by the current "
            f"user: {', '.join(deployment_names)} !"
        )

    if not yes:
        confirmation = cli_utils.confirmation(
            "Are you sure you want to continue?"
        )
        if not confirmation:
            cli_utils.declare("Deployment deletion canceled.")
            return

    for deployment in deployments:
        with console.status(f"Deleting deployment '{deployment.name}'...\n"):
            try:
                Client().delete_deployment(
                    name_id_or_prefix=deployment.id,
                    force=force,
                    timeout=timeout,
                )
                cli_utils.declare(f"Deleted deployment '{deployment.name}'.")
            except KeyError as e:
                error_message = (
                    f"Failed to delete deployment '{deployment.name}': "
                    f"{str(e)}"
                )
                if all or mine and ignore_errors:
                    cli_utils.warning(error_message)
                else:
                    cli_utils.error(error_message)


@deployment.command("refresh", help="Refresh the status of a deployment.")
@click.argument("deployment_name_or_id", type=str, required=True)
def refresh_deployment(
    deployment_name_or_id: str,
) -> None:
    """Refresh the status of a deployment.

    Args:
        deployment_name_or_id: The name or ID of the deployment to refresh.
    """
    try:
        deployment = Client().refresh_deployment(
            name_id_or_prefix=deployment_name_or_id
        )

    except KeyError as e:
        cli_utils.exception(e)
    else:
        cli_utils.pretty_print_deployment(deployment, show_secret=True)


@deployment.command(
    "invoke",
    context_settings={"ignore_unknown_options": True},
    help="Invoke a deployment with arguments.",
)
@click.argument("deployment_name_or_id", type=str, required=True)
@click.option(
    "--timeout",
    "-t",
    "timeout",
    type=int,
    required=False,
    default=None,
    help="Maximum time in seconds to wait for the deployment to be invoked.",
)
@click.argument("args", nargs=-1, type=click.UNPROCESSED)
def invoke_deployment(
    deployment_name_or_id: str,
    args: List[str],
    timeout: Optional[int] = None,
) -> None:
    """Call a deployment with arguments.

    Args:
        deployment_name_or_id: The name or ID of the deployment to call.
        args: The arguments to pass to the deployment call.
        timeout: The maximum time in seconds to wait for the deployment
            to be invoked.
    """
    from zenml.deployers.utils import invoke_deployment

    # Parse the given args
    args = list(args)
    args.append(deployment_name_or_id)

    name_or_id, parsed_args = cli_utils.parse_name_and_extra_arguments(
        args,
        expand_args=True,
        name_mandatory=True,
    )
    assert name_or_id is not None

    try:
        response = invoke_deployment(
            deployment_name_or_id=name_or_id,
            timeout=timeout or 300,  # 5 minute timeout
            project=None,
            **parsed_args,
        )
    except DeploymentInvalidParametersError as e:
        cli_utils.error(
            f"Invalid parameters for deployment '{name_or_id}': \n"
            f"{str(e)}\n\n"
            f"Hint: run 'zenml deployment describe --schema {name_or_id}' "
            "to inspect the deployment schema."
        )
    except KeyError as e:
        cli_utils.error(
            str(e)
            + "\n"
            + f"Hint: run [green]`zenml deployment logs {name_or_id}`[/green] "
            "to inspect the deployment logs."
        )
    else:
        cli_utils.declare(f"Invoked deployment '{name_or_id}' with response:")
        print(json.dumps(response, indent=2))
        if isinstance(response, dict) and not response.get("success", True):
            cli_utils.declare(
                f"Hint: run [green]`zenml deployment logs {name_or_id}`[/green] "
                "to inspect the deployment logs."
            )


@deployment.command("logs", help="Get the logs of a deployment.")
@click.argument("deployment_name_or_id", type=str, required=True)
@click.option(
    "--follow",
    "-f",
    is_flag=True,
    default=False,
    help="Follow the logs.",
)
@click.option(
    "--tail",
    "-t",
    type=int,
    default=None,
    help="The number of lines to show from the end of the logs.",
)
def log_deployment(
    deployment_name_or_id: str,
    follow: bool = False,
    tail: Optional[int] = None,
) -> None:
    """Get the logs of a deployment.

    Args:
        deployment_name_or_id: The name or ID of the deployment to get the logs of.
        follow: If True, follow the logs.
        tail: The number of lines to show from the end of the logs. If None,
            show all logs.
    """
    try:
        logs = Client().get_deployment_logs(
            name_id_or_prefix=deployment_name_or_id,
            follow=follow,
            tail=tail,
        )
    except KeyError as e:
        cli_utils.exception(e)
    else:
        for log in logs:
            print(log)
