#  Copyright (c) ZenML GmbH 2021. 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.
"""Functionality to install or uninstall ZenML integrations via the CLI."""

import os
import subprocess
import sys
from typing import Optional, Tuple

import click
from rich.progress import track

from zenml.cli.cli import TagGroup, cli
from zenml.cli.utils import (
    confirmation,
    declare,
    error,
    exception,
    format_integration_list,
    install_packages,
    print_table,
    title,
    uninstall_package,
    warning,
)
from zenml.console import console
from zenml.enums import CliCategories
from zenml.logger import get_logger

_WHYLOGS_INTEGRATION_WARNING = (
    "WhyLabs was acquired by Apple and the hosted WhyLabs platform is being "
    "discontinued. The whylogs library remains open source and the WhyLabs "
    "platform is now available as OSS at "
    "https://github.com/whylabs/whylabs-oss, but hosted functionality may stop "
    "working. Plan accordingly before continuing with the whylogs integration."
)

logger = get_logger(__name__)


@cli.group(
    cls=TagGroup,
    tag=CliCategories.INTEGRATIONS,
)
def integration() -> None:
    """Interact with external integrations."""


@integration.command(name="list", help="List the available integrations.")
def list_integrations() -> None:
    """List all available integrations with their installation status."""
    from zenml.integrations.registry import integration_registry

    formatted_table = format_integration_list(
        sorted(list(integration_registry.integrations.items()))
    )
    print_table(formatted_table)
    warning(
        "\n" + "To install the dependencies of a specific integration, type: "
    )
    warning("zenml integration install INTEGRATION_NAME")


@integration.command(
    name="requirements", help="List all requirements for an integration."
)
@click.argument("integration_name", required=False, default=None)
def get_requirements(integration_name: Optional[str] = None) -> None:
    """List all requirements for the chosen integration.

    Args:
        integration_name: The name of the integration to list the requirements
            for.
    """
    from zenml.integrations.registry import integration_registry

    try:
        requirements = integration_registry.select_integration_requirements(
            integration_name
        )
    except KeyError as e:
        exception(e)
    else:
        if requirements:
            title(
                f"Requirements for {integration_name or 'all integrations'}:\n"
            )
            declare(f"{requirements}")
            warning(
                "\n" + "To install the dependencies of a "
                "specific integration, type: "
            )
            warning("zenml integration install INTEGRATION_NAME")


@integration.command(
    name="export-requirements", help="Export the integration requirements."
)
@click.argument("integrations", nargs=-1, required=False)
@click.option(
    "--ignore-integration",
    "-i",
    multiple=True,
    help="List of integrations to ignore explicitly.",
)
@click.option(
    "--output-file",
    "-o",
    "output_file",
    type=str,
    required=False,
    help="File to which to export the integration requirements. If not "
    "provided, the requirements will be printed to stdout instead.",
)
@click.option(
    "--overwrite",
    "-ov",
    "overwrite",
    type=bool,
    required=False,
    is_flag=True,
    help="Overwrite the output file if it already exists. This option is "
    "only valid if the output file is provided.",
)
@click.option(
    "--installed-only",
    "installed_only",
    is_flag=True,
    default=False,
    help="Only export requirements for integrations installed in your current "
    "environment. This can not be specified when also providing explicit "
    "integrations.",
)
@click.option(
    "--poetry",
    "poetry",
    is_flag=True,
    default=False,
    help="Add the exported requirements to your current Poetry project.",
)
def export_requirements(
    integrations: Tuple[str],
    ignore_integration: Tuple[str],
    output_file: Optional[str] = None,
    overwrite: bool = False,
    installed_only: bool = False,
    poetry: bool = False,
) -> None:
    """Exports integration requirements so they can be installed using pip.

    Args:
        integrations: The name of the integration to install the requirements
            for.
        ignore_integration: List of integrations to ignore explicitly.
        output_file: Optional path to the requirements output file.
        overwrite: Overwrite the output file if it already exists. This option
            is only valid if the output file is provided.
        installed_only: Only export requirements for integrations installed in
            your current environment. This can not be specified when also
            providing explicit integrations.
        poetry: Add the exported requirements to your current Poetry project.
    """
    from zenml.integrations.registry import integration_registry

    if installed_only and integrations:
        error(
            "You can either provide specific integrations or export only "
            "requirements for integrations installed in your local "
            "environment, not both."
        )

    if poetry and output_file:
        error(
            "You can either specify an output file or add the requirements to "
            "the Poetry project, not both."
        )

    all_integrations = set(integration_registry.integrations.keys())

    if integrations:
        integrations_to_export = set(integrations)
    elif installed_only:
        integrations_to_export = set(
            integration_registry.get_installed_integrations()
        )
    else:
        integrations_to_export = all_integrations

    for i in ignore_integration:
        try:
            integrations_to_export.remove(i)
        except KeyError:
            if i not in all_integrations:
                error(
                    f"Integration {i} does not exist. Available integrations: "
                    f"{all_integrations}"
                )

    if "whylogs" in integrations_to_export:
        warning(_WHYLOGS_INTEGRATION_WARNING)

    requirements = []
    for integration_name in integrations_to_export:
        try:
            requirements += (
                integration_registry.select_integration_requirements(
                    integration_name
                )
            )
        except KeyError:
            error(f"Unable to find integration '{integration_name}'.")

    if output_file:
        try:
            with open(output_file, "x") as f:
                f.write("\n".join(requirements))
        except FileExistsError:
            if overwrite or confirmation(
                "A file already exists at the specified path. "
                "Would you like to overwrite it?"
            ):
                with open(output_file, "w") as f:
                    f.write("\n".join(requirements))
        declare(f"Requirements exported to {output_file}.")
    if poetry:
        res = os.popen("poetry env list").read()
        envs = [
            env
            for env in res.split("\n")
            if env.lower().find("(activated)") > 0
        ]
        if len(envs) == 0:
            error(
                "No activated Poetry environment found. Please activate one "
                "and try again."
            )
        else:
            # Use subprocess.run with shell=False to avoid command injection
            args = ["poetry", "add"] + requirements
            subprocess.run(args, check=True)
            declare(
                f"Requirements added to `{envs[0]}` environment in Poetry."
            )
    else:
        click.echo(" ".join(requirements), nl=False)


@integration.command(
    help="Install the required packages for the integration of choice."
)
@click.argument("integrations", nargs=-1, required=False)
@click.option(
    "--ignore-integration",
    "-i",
    multiple=True,
    help="Integrations to ignore explicitly (passed in separately).",
)
@click.option(
    "--yes",
    "-y",
    "force",
    is_flag=True,
    help="Force the installation of the required packages. This will skip the "
    "confirmation step and reinstall existing packages as well",
)
@click.option(
    "--uv",
    "uv",
    is_flag=True,
    help="Experimental: Use uv for package installation.",
    default=False,
)
def install(
    integrations: Tuple[str],
    ignore_integration: Tuple[str],
    force: bool = False,
    uv: bool = False,
) -> None:
    """Installs the required packages for a given integration.

    If no integration is specified all required packages for all integrations
    are installed using pip or uv.

    Args:
        integrations: The name of the integration to install the requirements
            for.
        ignore_integration: Integrations to ignore explicitly (passed in
            separately).
        force: Force the installation of the required packages.
        uv: Use uv for package installation (experimental).
    """
    from zenml.cli.utils import is_pip_installed, is_uv_installed
    from zenml.integrations.registry import integration_registry

    if uv and not is_uv_installed():
        error(
            "UV is not installed but the uv flag was passed in. Please install "
            "uv or remove the uv flag."
        )

    if not uv and not is_pip_installed():
        error(
            "Pip is not installed. Please install pip or use the uv flag "
            "(--uv) for package installation."
        )

    if not integrations:
        # no integrations specified, use all registered integrations
        integration_set = set(integration_registry.integrations.keys())

        for i in ignore_integration:
            try:
                integration_set.remove(i)
            except KeyError:
                error(
                    f"Integration {i} does not exist. Available integrations: "
                    f"{list(integration_registry.integrations.keys())}"
                )
    else:
        integration_set = set(integrations)

    if "whylogs" in integration_set:
        warning(_WHYLOGS_INTEGRATION_WARNING)

    if sys.version_info.minor == 12 and "tensorflow" in integration_set:
        warning(
            "The TensorFlow integration is not yet compatible with Python "
            "3.12, thus its installation is skipped. Consider using a "
            "different version of Python and stay in touch for further updates."
        )
        integration_set.remove("tensorflow")

    if sys.version_info.minor == 12 and "deepchecks" in integration_set:
        warning(
            "The Deepchecks integration is not yet compatible with Python "
            "3.12, thus its installation is skipped. Consider using a "
            "different version of Python and stay in touch for further updates."
        )
        integration_set.remove("deepchecks")

    requirements = []
    integrations_to_install = []
    for integration_name in integration_set:
        try:
            if force or not integration_registry.is_installed(
                integration_name
            ):
                requirements += (
                    integration_registry.select_integration_requirements(
                        integration_name
                    )
                )
                integrations_to_install.append(integration_name)
            else:
                declare(
                    f"All required packages for integration "
                    f"'{integration_name}' are already installed."
                )
        except KeyError:
            error(f"Unable to find integration '{integration_name}'.")

    if requirements and (
        force
        or confirmation(
            "Are you sure you want to install the following "
            "packages to the current environment?\n"
            f"{requirements}"
        )
    ):
        with console.status("Installing integrations..."):
            install_packages(requirements, use_uv=uv)


@integration.command(
    help="Uninstall the required packages for the integration of choice."
)
@click.argument("integrations", nargs=-1, required=False)
@click.option(
    "--yes",
    "-y",
    "force",
    is_flag=True,
    help="Force the uninstallation of the required packages. This will skip "
    "the confirmation step",
)
@click.option(
    "--uv",
    "uv",
    is_flag=True,
    help="Experimental: Use uv for package uninstallation.",
    default=False,
)
def uninstall(
    integrations: Tuple[str], force: bool = False, uv: bool = False
) -> None:
    """Uninstalls the required packages for a given integration.

    If no integration is specified all required packages for all integrations
    are uninstalled using pip or uv.

    Args:
        integrations: The name of the integration to uninstall the requirements
            for.
        force: Force the uninstallation of the required packages.
        uv: Use uv for package uninstallation (experimental).
    """
    from zenml.cli.utils import is_pip_installed, is_uv_installed
    from zenml.integrations.registry import integration_registry

    if uv and not is_uv_installed():
        error("Package `uv` is not installed. Please install it and retry.")

    if not uv and not is_pip_installed():
        error(
            "Pip is not installed. Please install pip or use the uv flag "
            "(--uv) for package installation."
        )

    if not integrations:
        # no integrations specified, use all registered integrations
        integrations = tuple(integration_registry.integrations.keys())

    requirements = []
    for integration_name in integrations:
        try:
            if integration_registry.is_installed(integration_name):
                requirements += (
                    integration_registry.select_uninstall_requirements(
                        integration_name
                    )
                )
            else:
                warning(
                    f"Requirements for integration '{integration_name}' "
                    f"already not installed."
                )
        except KeyError:
            warning(f"Unable to find integration '{integration_name}'.")

    if requirements and (
        force
        or confirmation(
            "Are you sure you want to uninstall the following "
            "packages from the current environment?\n"
            f"{requirements}"
        )
    ):
        for n in track(
            range(len(requirements)),
            description="Uninstalling integrations...",
        ):
            uninstall_package(requirements[n], use_uv=uv)


@integration.command(
    help="Upgrade the required packages for the integration of choice."
)
@click.argument("integrations", nargs=-1, required=False)
@click.option(
    "--yes",
    "-y",
    "force",
    is_flag=True,
    help="Force the upgrade of the required packages. This will skip the "
    "confirmation step and re-upgrade existing packages as well",
)
@click.option(
    "--uv",
    "uv",
    is_flag=True,
    help="Experimental: Use uv for package upgrade.",
    default=False,
)
def upgrade(
    integrations: Tuple[str],
    force: bool = False,
    uv: bool = False,
) -> None:
    """Upgrade the required packages for a given integration.

    If no integration is specified all required packages for all integrations
    are installed using pip or uv.

    Args:
        integrations: The name of the integration to install the requirements
            for.
        force: Force the installation of the required packages.
        uv: Use uv for package installation (experimental).
    """
    from zenml.cli.utils import is_pip_installed, is_uv_installed
    from zenml.integrations.registry import integration_registry

    if uv and not is_uv_installed():
        error("Package `uv` is not installed. Please install it and retry.")

    if not uv and not is_pip_installed():
        error(
            "Pip is not installed. Please install pip or use the uv flag "
            "(--uv) for package installation."
        )

    if not integrations:
        # no integrations specified, use all registered integrations
        integrations = set(integration_registry.integrations.keys())

    requirements = []
    integrations_to_install = []
    for integration_name in integrations:
        try:
            if integration_registry.is_installed(integration_name):
                requirements += (
                    integration_registry.select_integration_requirements(
                        integration_name
                    )
                )
                integrations_to_install.append(integration_name)
            else:
                declare(
                    f"None of the required packages for integration "
                    f"'{integration_name}' are installed."
                )
        except KeyError:
            warning(f"Unable to find integration '{integration_name}'.")

    if requirements and (
        force
        or confirmation(
            f"Are you sure you want to upgrade the following "
            "packages to the current environment?\n"
            f"{requirements}"
        )
    ):
        with console.status("Upgrading integrations..."):
            install_packages(requirements, upgrade=True, use_uv=uv)
