from json import dumps as json_dumps
import pathlib
from typing import Dict, Optional, Tuple

import click
from typing_extensions import Literal

from anyscale.cli_logger import BlockLogger
from anyscale.controllers.service_controller import ServiceController
import anyscale.service
from anyscale.service import ServiceConfig, ServiceStatus
from anyscale.util import validate_non_negative_arg


log = BlockLogger()  # CLI Logger


@click.group("service")
def service_cli():
    pass


def _convert_arguments_to_dict(args: Tuple[str]) -> Dict[str, str]:
    """Convert args of the form "key=val" into a dictionary of {key: val}.

    NOTE(edoakes): this mimics the functionality of the `serve run` CLI and should be
    kept in sync with it.
    """
    args_dict = {}
    for arg in args:
        split = arg.split("=")
        if len(split) != 2:
            raise click.ClickException(
                f"Invalid application argument '{arg}', "
                "must be of the form '<key>=<val>'."
            )

        args_dict[split[0]] = split[1]

    return args_dict


@service_cli.command(
    name="deploy",
    short_help="Deploy or update a service.",
    help=(
        "Deploy or update a service.\n\n"
        "If no service with the provided name is running, one will be created, "
        "else the existing service will be updated.\n\n"
        "The target of this command can either be an import path to a Ray Serve "
        "application or a path to a service config file.\n\n"
        "Command-line flags take precedence over values in the config file."
    ),
    hidden=True,
)
@click.argument("import_path_or_config", type=str)
@click.argument("arguments", nargs=-1, required=False)
@click.option(
    "-n",
    "--name",
    required=False,
    default=None,
    type=str,
    help="Unique name for the service.",
)
@click.option(
    "--image",
    required=False,
    default=None,
    type=str,
    help="Container image to use for the service. This cannot be changed when using --in-place. When running in a workspace, this defaults to the image of the workspace.",
)
@click.option(
    "--working-dir",
    required=False,
    default=None,
    type=str,
    help="Path to a local directory that will be the working directory for the service. The files in the directory will be automatically uploaded to cloud storage. When running in a workspace, this defaults to the current working directory.",
)
@click.option(
    "-r",
    "--requirements",
    required=False,
    default=None,
    type=str,
    help="Path to a requirements.txt file containing dependencies for the service. These will be installed on top of the image. When running in a workspace, this defaults to the workspace dependencies.",
)
@click.option(
    "-i",
    "--in-place",
    is_flag=True,
    show_default=True,
    default=False,
    help="Perform an in-place upgrade without starting a new cluster. This can be used for faster iteration during development but is *not* currently recommended for production deploys.",
)
@click.option(
    "--canary-percent",
    required=False,
    default=None,
    type=int,
    help="The percentage of traffic to send to the canary version of the service (0-100). This can be used to manually shift traffic toward (or away from) the canary version. If not provided, traffic will be shifted incrementally toward the canary version until it reaches 100. Not supported when using --in-place.",
)
@click.option(
    "--max-surge-percent",
    required=False,
    default=None,
    type=int,
    help="Amount of excess capacity allowed to be used while updating the service (0-100). Defaults to 100. Not supported when using --in-place.",
)
def deploy(  # noqa: PLR0913
    import_path_or_config: str,
    arguments: Tuple[str],
    name: Optional[str],
    image: Optional[str],
    working_dir: Optional[str],
    requirements: Optional[str],
    in_place: bool,
    canary_percent: Optional[int],
    max_surge_percent: Optional[int],
):
    if import_path_or_config.endswith(".yaml"):
        if not pathlib.Path(import_path_or_config).is_file():
            raise click.ClickException(
                f"Config file '{import_path_or_config}' not found."
            )

        if len(arguments) > 0:
            raise click.ClickException(
                "Application arguments can't be passed when deploying from a config file. "
                "Include them in the config file instead."
            )

        config = ServiceConfig.from_yaml(import_path_or_config)

    else:
        config = ServiceConfig(
            import_path=import_path_or_config,
            arguments=_convert_arguments_to_dict(arguments) or None,
        )

    anyscale.service.deploy(
        config.options(
            name=name, image=image, working_dir=working_dir, requirements=requirements,
        ),
        in_place=in_place,
        canary_percent=canary_percent,
        max_surge_percent=max_surge_percent,
    )


@service_cli.command(
    name="status", help="Get the status of a service.", hidden=True,
)
@click.argument("name", type=str)
@click.option(
    "--json",
    "-j",
    is_flag=True,
    show_default=True,
    default=False,
    help="Output the status in a structured JSON format.",
)
def status(name: str, json: bool):
    status: ServiceStatus = anyscale.service.status(name)
    if json:
        print(json_dumps(status.to_dict()))
    else:
        print(status)


@service_cli.command(
    name="rollout", help="Roll out a service.",
)
@click.option(
    "-f",
    "--service-config-file",
    required=True,
    help="The path of the service configuration file.",
)
@click.option("--name", "-n", required=False, default=None, help="Name of service.")
@click.option("--version", required=False, default=None, help="Version of service.")
@click.option(
    "--max-surge-percent",
    required=False,
    default=None,
    type=int,
    help="Max amount of excess capacity allocated during the rollout (0-100).",
)
@click.option(
    "--canary-percent",
    required=False,
    default=None,
    type=int,
    help="The percentage of traffic to send to the canary version of the service.",
)
@click.option(
    "--rollout-strategy",
    required=False,
    default=None,
    type=click.Choice(["ROLLOUT", "IN_PLACE"], case_sensitive=False),
    help="Strategy for rollout.",
)
@click.option(
    "-i",
    "--in-place",
    "in_place",
    is_flag=True,
    show_default=True,
    default=False,
    help="Alias for `--rollout-strategy=IN_PLACE`.",
)
@click.option(
    "--no-auto-complete-rollout",
    is_flag=True,
    show_default=True,
    default=False,
    help="Do not complete the rollout (terminate the existing version cluster) after the canary version reaches 100%",
)
def rollout(  # noqa: PLR0913
    service_config_file: str,
    name: Optional[str],
    version: Optional[str],
    max_surge_percent: Optional[int],
    canary_percent: Optional[int],
    rollout_strategy: Optional[Literal["ROLLOUT", "IN_PLACE"]],
    in_place: bool,
    no_auto_complete_rollout: bool,
):
    """Start or update a service rollout to a new version."""
    if in_place:
        if rollout_strategy is not None:
            raise click.ClickException(
                "Only one of `--in-place/-i` and `--rollout-strategy` can be provided."
            )
        rollout_strategy = "IN_PLACE"

    service_controller = ServiceController()
    service_config = service_controller.read_service_config_file(service_config_file)
    service_controller.rollout(
        service_config,
        name=name,
        version=version,
        max_surge_percent=max_surge_percent,
        canary_percent=canary_percent,
        rollout_strategy=rollout_strategy,
        auto_complete_rollout=not no_auto_complete_rollout,
    )


@service_cli.command(name="list", help="Display information about services.")
@click.option(
    "--name", "-n", required=False, default=None, help="Filter by service name."
)
@click.option(
    "--service-id", "--id", required=False, default=None, help="Filter by service id."
)
@click.option(
    "--project-id", required=False, default=None, help="Filter by project id."
)
@click.option(
    "--created-by-me",
    help="List services created by me only.",
    is_flag=True,
    default=False,
)
@click.option(
    "--max-items",
    required=False,
    default=10,
    type=int,
    help="Max items to show in list.",
    callback=validate_non_negative_arg,
)
def list(  # noqa: A001
    name: Optional[str],
    service_id: Optional[str],
    project_id: Optional[str],
    created_by_me: bool,
    max_items: int,
):
    """List services based on the provided filters.

    This returns both v1 and v2 services.
    """
    service_controller = ServiceController()
    service_controller.list(
        name=name,
        service_id=service_id,
        project_id=project_id,
        created_by_me=created_by_me,
        max_items=max_items,
    )


@service_cli.command(name="rollback", help="Roll back a service.")
@click.option(
    "--service-id", "--id", default=None, help="ID of service.",
)
@click.option("--name", "-n", required=False, default=None, help="Name of service.")
@click.option("--project-id", required=False, help="Filter by project id.")
@click.option(
    "--service-config-file", "-f", help="The path of the service configuration file.",
)
@click.option(
    "--max-surge-percent",
    required=False,
    default=None,
    type=int,
    help="Max amount of excess capacity allocated during the rollback (0-100).",
)
def rollback(
    service_id: Optional[str],
    name: Optional[str],
    project_id: Optional[str],
    service_config_file: Optional[str],
    max_surge_percent: Optional[int],
):
    """Perform a rollback for a service that is currently in a rollout.

    This *only* applies to v2 services.
    """
    service_controller = ServiceController()
    service_id = service_controller.get_service_id(
        service_id=service_id,
        service_name=name,
        service_config_file=service_config_file,
        project_id=project_id,
    )
    service_controller.rollback(service_id, max_surge_percent)


@service_cli.command(name="terminate", help="Terminate a service.")
@click.option(
    "--service-id", "--id", required=False, help="ID of service.",
)
@click.option("--name", "-n", required=False, help="Name of service.")
@click.option("--project-id", required=False, help="Filter by project id.")
@click.option(
    "--service-config-file", "-f", help="The path of the service configuration file.",
)
def terminate(
    service_id: Optional[str],
    name: Optional[str],
    project_id: Optional[str],
    service_config_file: Optional[str],
):
    """Terminate a service.

    This applies to both v1 and v2 services.
    """
    service_controller = ServiceController()
    service_id = service_controller.get_service_id(
        service_id=service_id,
        service_name=name,
        service_config_file=service_config_file,
        project_id=project_id,
    )
    service_controller.terminate(service_id)
