import os
from typing import Any, Dict, List, Optional, Union

import click
import tabulate
import yaml

from anyscale.cli_logger import BlockLogger
from anyscale.client.openapi_client import ComputeTemplate, ProductionJob
from anyscale.client.openapi_client.models.decorated_list_service_api_model import (
    DecoratedListServiceAPIModel,
)
from anyscale.client.openapi_client.models.decorated_production_job import (
    DecoratedProductionJob,
)
from anyscale.client.openapi_client.models.decorated_production_service_v2_api_model import (
    DecoratedProductionServiceV2APIModel,
)
from anyscale.client.openapi_client.models.ha_job_states import HaJobStates
from anyscale.client.openapi_client.models.service_event_current_state import (
    ServiceEventCurrentState,
)
from anyscale.client.openapi_client.models.service_goal_states import ServiceGoalStates
from anyscale.controllers.base_controller import BaseController
from anyscale.controllers.job_controller import JobController
from anyscale.formatters.service_formatter import (
    format_service_config,
    format_service_config_v2,
)
from anyscale.models.service_model import ServiceConfig
from anyscale.project import infer_project_id
from anyscale.tables import ServicesTable
from anyscale.util import get_endpoint, is_anyscale_workspace
from anyscale.utils.runtime_env import (
    is_dir_remote_uri,
    override_runtime_env_for_local_working_dir,
    upload_and_rewrite_working_dir,
    upload_and_rewrite_workspace_working_dir,
)


def _is_v2_service_id(service_id: str) -> bool:
    return service_id.startswith("service2")


class ServiceController(BaseController):
    def __init__(
        self, log: Optional[BlockLogger] = None, initialize_auth_api_client: bool = True
    ):
        if log is None:
            log = BlockLogger()

        super().__init__(initialize_auth_api_client=initialize_auth_api_client)
        self.log = log
        self.job_controller = JobController(
            initialize_auth_api_client=initialize_auth_api_client
        )

    def _get_services_by_name(
        self,
        *,
        name: Optional[str] = None,
        project_id: Optional[str] = None,
        created_by_me: bool,
        max_items: int,
    ) -> List[DecoratedListServiceAPIModel]:
        """Makes an API call to get all services matching the provided filters.

        This returns both v1 and v2 services.
        """
        creator_id = None
        if created_by_me:
            user_info_response = self.api_client.get_user_info_api_v2_userinfo_get()
            creator_id = user_info_response.result.id

        services_list = []
        resp = self.api_client.list_services_api_v2_services_v2_get(
            project_id=project_id, name=name, creator_id=creator_id, count=10,
        )
        services_list.extend(resp.results)
        paging_token = resp.metadata.next_paging_token
        while paging_token is not None and len(services_list) < max_items:
            resp = self.api_client.list_services_api_v2_services_v2_get(
                project_id=project_id,
                name=name,
                creator_id=creator_id,
                count=10,
                paging_token=paging_token,
            )
            services_list.extend(resp.results)
            paging_token = resp.metadata.next_paging_token

        return services_list[:max_items]

    def _get_service_id_from_name(
        self, service_name: str, project_id: Optional[str]
    ) -> str:
        """Get the ID for a service by name.

        If project_id is specified, filter to that project, else don't filter on project_id and
        instead error if there are multiple services with the name.

        Raises an exception if there are zero or multiple services with the given name.
        """
        results = self._get_services_by_name(
            name=service_name, project_id=project_id, created_by_me=False, max_items=10
        )

        if len(results) == 0:
            raise click.ClickException(
                f"No service with name '{service_name}' was found. "
                "Please verify that this service exists and you have access to it."
            )
        elif len(results) > 1:
            raise click.ClickException(
                f"There are multiple services with name '{service_name}'. "
                "Please filter using --project-id or specify the --service-id instead. "
                f"Services found: \n{ServicesTable(results)}"
            )

        return results[0].id

    def get_service_id(
        self,
        *,
        service_id: Optional[str] = None,
        service_name: Optional[str] = None,
        service_config_file: Optional[str] = None,
        project_id: Optional[str] = None,
    ) -> str:
        """Get the service ID given the ID, name, or config file.

        This is a utility used by multiple CLI commands to standardize mapping these options
        to a service_id.

        The precedence is: service_id > service_name > service_config_file.

        If the project_id is specified directly or via the service_config_file it will be used
        to filter the query. Else we will try to find the service across all projects and error
        if there are multiple.
        """
        if service_id is not None:
            if service_name is not None or service_config_file is not None:
                raise click.ClickException(
                    "Only one of service ID, name, or config file should be specified."
                )
        elif service_name is not None:
            if service_config_file is not None:
                raise click.ClickException(
                    "Only one of service ID, name, or config file should be specified."
                )
            service_id = self._get_service_id_from_name(service_name, project_id)
        elif service_config_file is not None:
            service_config: ServiceConfig = self._read_service_config_file(
                service_config_file
            )
            # Allow the passed project_id to override the one in the file.
            project_id = project_id or service_config.project_id
            service_id = self._get_service_id_from_name(service_config.name, project_id)
        else:
            raise click.ClickException(
                "Service ID, name, or config file must be specified."
            )

        return service_id

    def deploy(
        self,
        service_config_file: str,
        name: Optional[str] = None,
        description: Optional[str] = None,
        healthcheck_url: Optional[str] = None,
        is_entrypoint_cmd: Optional[bool] = False,
        entrypoint: Optional[List[str]] = None,
    ) -> None:
        entrypoint = entrypoint or []
        if is_anyscale_workspace() and is_entrypoint_cmd:
            # If is_entrypoint_cmd is true, service_config_file is expected to be "python"
            # Therefore, we need to combine args to create the final entrypoint
            entrypoint = [service_config_file, *entrypoint]
            config = self.generate_config_from_entrypoint(
                entrypoint,
                name=name,
                description=description,
                healthcheck_url=healthcheck_url,
            )
            project_id = infer_project_id(
                self.anyscale_api_client,
                self.api_client,
                self.log,
                project_id=config.project_id,
                cluster_compute_id=config.compute_config_id,
                cluster_compute=config.compute_config,
                cloud=config.cloud,
            )

            config.runtime_env = override_runtime_env_for_local_working_dir(
                config.runtime_env
            )
            config.project_id = project_id

            self.deploy_from_config(config)
        elif len(entrypoint) == 0:
            # Assume that service_config_file is a file and submit it.
            config = self.generate_config_from_file(
                service_config_file,
                name=name,
                description=description,
                healthcheck_url=healthcheck_url,
            )

            project_id = infer_project_id(
                self.anyscale_api_client,
                self.api_client,
                self.log,
                project_id=config.project_id,
                cluster_compute_id=config.compute_config_id,
                cluster_compute=config.compute_config,
                cloud=config.cloud,
            )

            config.project_id = project_id
            self.deploy_from_config(config)
        elif len(entrypoint) != 0:
            msg = (
                "Within an Anyscale Workspace, `anyscale service deploy` takes either a file, or a command. To submit a command, use `anyscale service deploy -- my command`."
                if is_anyscale_workspace()
                else "`anyscale service deploy` takes one argument, a YAML file configuration. Please use `anyscale service deploy my_file`."
            )
            raise click.ClickException(msg)

    def generate_config_from_entrypoint(
        self,
        entrypoint: List[str],
        name: Optional[str] = None,
        description: Optional[str] = None,
        healthcheck_url: Optional[str] = "/healthcheck",
    ) -> ServiceConfig:
        config_dict = {
            "entrypoint": " ".join(entrypoint),
            "name": name,
            "description": description,
            "healthcheck_url": healthcheck_url,
        }
        return self._populate_service_config(config_dict)

    def generate_config_from_file(
        self,
        service_config_file,
        *,
        name: Optional[str] = None,
        description: Optional[str] = None,
        healthcheck_url: Optional[str] = None,
        version: Optional[str] = None,
        canary_percent: Optional[int] = None,
        rollout_strategy: Optional[str] = None,
    ) -> ServiceConfig:
        """
        Given the input file path and user overrides, this method
        constructs the expected service configuration.

        Note that if the command is run from a workspace, some values
        may be overwritten.
        """
        service_config: ServiceConfig = self._read_service_config_file(
            service_config_file
        )
        if name:
            service_config.name = name

        if description:
            service_config.description = description

        if healthcheck_url:
            service_config.healthcheck_url = healthcheck_url

        if version:
            service_config.version = version

        if canary_percent is not None:
            service_config.canary_percent = canary_percent

        if rollout_strategy:
            service_config.rollout_strategy = rollout_strategy

        if service_config.rollout_strategy:
            service_config.rollout_strategy = service_config.rollout_strategy.upper()

        self._override_configs_in_workspaces(service_config)

        if service_config.runtime_env:
            # Service v1 specific override since Service v2's runtime_env
            # is under service_config.ray_serve_config["runtime_env"]
            service_config.runtime_env = upload_and_rewrite_working_dir(
                service_config.runtime_env
            )

        return service_config

    def _override_configs_in_workspaces(self, service_config: ServiceConfig) -> None:
        """
        We override some service configs if launched in workspaces.

        Please note that this method mutates service_config in-place.
        """
        if (
            is_anyscale_workspace()
            and "ANYSCALE_SESSION_ID" in os.environ
            and service_config.ray_serve_config
        ):
            existing_runtime_env = service_config.ray_serve_config.get(
                "runtime_env", {}
            )

            # If deploying service v2 from workspaces and no upload_path is specified,
            # we back up the current workspace content into S3
            working_dir = existing_runtime_env.get("working_dir", "")
            upload_path = existing_runtime_env.get("upload_path")
            if not is_dir_remote_uri(working_dir) and upload_path is None:
                workspace_id = os.environ["ANYSCALE_EXPERIMENTAL_WORKSPACE_ID"]
                cluster_id = os.environ["ANYSCALE_SESSION_ID"]
                new_runtime_env = upload_and_rewrite_workspace_working_dir(
                    api_client=self.api_client,
                    existing_runtime_env=existing_runtime_env,
                    workspace_id=workspace_id,
                    cluster_id=cluster_id,
                    log=self.log,
                )
                service_config.ray_serve_config["runtime_env"] = new_runtime_env

    def deploy_from_config(self, service_config: ServiceConfig):
        apply_service_config = format_service_config(service_config)
        # TODO: shawnp. Change to the new Services API
        # once the new endpoint works for both Service v1 and v2.
        create_production_service = apply_service_config.service_v1_config
        if create_production_service:
            service = self.api_client.apply_service_api_v2_decorated_ha_jobs_apply_service_put(
                create_production_service
            ).result
        else:
            # TODO: shawnp. This condition should never occur
            raise RuntimeError("The Service configuration is missing")
        service_id = service.id
        current_state = self._format_service_state(service.state.current_state)

        maximum_uptime_minutes = self._get_maximum_uptime_minutes(service)
        self.log.info(
            f"Maximum uptime is {self._get_maximum_uptime_output(maximum_uptime_minutes)} for clusters launched by this service."
            f"{self._get_additional_log_if_maximum_uptime_enabled(maximum_uptime_minutes)}"
        )

        self.log.info(
            f"Service {service_id} has been deployed. Current state of service: {current_state}."
        )
        self.log.info(
            f"Query the status of the service with `anyscale service list --service-id {service.id}`."
        )
        self.log.info(
            f'View the service in the UI at {get_endpoint(f"/services/{service.id}")}.'
        )

    def _get_maximum_uptime_minutes(self, service: ProductionJob) -> Optional[int]:
        compute_config: ComputeTemplate = self.api_client.get_compute_template_api_v2_compute_templates_template_id_get(
            service.config.compute_config_id
        ).result
        return compute_config.config.maximum_uptime_minutes

    def _get_maximum_uptime_output(self, maximum_uptime_minutes: Optional[int]) -> str:
        if maximum_uptime_minutes and maximum_uptime_minutes > 0:
            return f"set to {maximum_uptime_minutes} minutes"
        return "disabled"

    def _get_additional_log_if_maximum_uptime_enabled(
        self, maximum_uptime_minutes: Optional[int]
    ) -> str:
        if maximum_uptime_minutes and maximum_uptime_minutes > 0:
            return " This may cause disruptions. To disable, update the compute config."
        return ""

    def _populate_service_config(self, config_dict: Dict[str, Any]) -> ServiceConfig:
        """
        If the service is deployed from the workspace, we populate values from the workspace
        if the config doesn't specify them, such as cluster env and compute config.
        """
        if (
            "ANYSCALE_EXPERIMENTAL_WORKSPACE_ID" in os.environ
            and "ANYSCALE_SESSION_ID" in os.environ
        ):
            cluster = self.anyscale_api_client.get_cluster(
                os.environ["ANYSCALE_SESSION_ID"]
            ).result
            # If the job configs are not specified, infer them from the workspace:
            if "build_id" not in config_dict and "cluster_env" not in config_dict:
                config_dict["build_id"] = cluster.cluster_environment_build_id
            if "project_id" not in config_dict:
                config_dict["project_id"] = cluster.project_id
            if (
                "compute_config" not in config_dict
                and "compute_config_id" not in config_dict
            ):
                config_dict["compute_config_id"] = cluster.cluster_compute_id
        service_config = ServiceConfig.parse_obj(config_dict)
        return service_config

    def _read_service_config_file(self, service_config_file: str) -> ServiceConfig:
        if not os.path.exists(service_config_file):
            raise click.ClickException(f"Config file {service_config_file} not found.")

        with open(service_config_file) as f:
            config_dict = yaml.safe_load(f)

        return self._populate_service_config(config_dict)

    def rollout(
        self,
        service_config_file: str,
        name: Optional[str] = None,
        version: Optional[str] = None,
        canary_percent: Optional[int] = None,
        rollout_strategy: Optional[str] = None,
    ):
        """
        Deploys a Service 2.0.
        """
        config = self.generate_config_from_file(
            service_config_file,
            name=name,
            version=version,
            canary_percent=canary_percent,
            rollout_strategy=rollout_strategy,
        )

        service_v2_config = format_service_config_v2(config)
        service = self.api_client.apply_service_v2_api_v2_services_v2_apply_put(
            service_v2_config
        ).result

        current_state = self._format_service_state(service.goal_state)
        self.log.info(
            f"Service {service.id} has been deployed. Service is transitioning towards: {current_state}."
        )
        self.log.info(
            f'View the service in the UI at {get_endpoint(f"/services/{service.id}")}.'
        )

    def list(
        self,
        *,
        name: Optional[str] = None,
        service_id: Optional[str] = None,
        project_id: Optional[str] = None,
        created_by_me: bool = False,
        max_items: int = 10,
    ) -> None:
        try:
            self._list_via_service_list_endpoint(
                name=name,
                service_id=service_id,
                project_id=project_id,
                created_by_me=created_by_me,
                max_items=max_items,
            )
        except Exception as e:  # noqa: BLE001
            if "Services v2 is not enabled for your organization." in str(e):
                # An exception will be thrown if the "services_v2_enable_all_features" flag is not enabled
                # Therefore, we default to showing only Service v1.
                # TODO: shawnp. Remove the exception logic once we enable service v2 flag for all customers
                self.job_controller.list(
                    name=name,
                    job_id=service_id,
                    project_id=project_id,
                    include_all_users=not created_by_me,
                    include_archived=False,
                    max_items=max_items,
                    is_service=True,
                )
            else:
                self.log.info(
                    "We are unable to get a list of Services."
                    "This is unexpected and please contact Anyscale support for more info."
                )

    def _list_via_service_list_endpoint(
        self,
        *,
        name: Optional[str] = None,
        service_id: Optional[str] = None,
        project_id: Optional[str] = None,
        created_by_me: bool = False,
        max_items: int,
    ):
        services = []
        if service_id:
            if _is_v2_service_id(service_id):
                service_v2_result: DecoratedProductionServiceV2APIModel = (
                    self.api_client.get_service_api_v2_services_v2_service_id_get(
                        service_id
                    ).result
                )
                services.append(
                    [
                        service_v2_result.name,
                        service_v2_result.id,
                        self._format_service_state(service_v2_result.current_state),
                        # TODO: change to email once https://github.com/anyscale/product/pull/18189 is merged
                        service_v2_result.creator_id,
                    ]
                )
            else:
                service_v1_result: DecoratedProductionJob = self.api_client.get_job_api_v2_decorated_ha_jobs_production_job_id_get(
                    service_id
                ).result
                services.append(
                    [
                        service_v1_result.name,
                        service_v1_result.id,
                        self._format_service_state(
                            service_v1_result.state.current_state
                        ),
                        service_v1_result.creator.email,
                    ]
                )

        else:
            services_data = self._get_services_by_name(
                name=name,
                project_id=project_id,
                created_by_me=created_by_me,
                max_items=max_items,
            )

            services = [
                [
                    service.name,
                    service.id,
                    self._format_service_state(service.current_state),
                    service.creator.email,
                ]
                for service in services_data
            ]

        table = tabulate.tabulate(
            services,
            headers=["NAME", "ID", "CURRENT STATE", "CREATOR"],
            tablefmt="plain",
        )
        self.log.info(f'View your Services in the UI at {get_endpoint("/services")}')
        self.log.info(f"Services:\n{table}")

    def _format_service_state(
        self, state: Union[ServiceEventCurrentState, HaJobStates, ServiceGoalStates]
    ):
        if state == "RUNNING":
            return "ACTIVE"
        return state

    def archive(self, service_id: str):
        if _is_v2_service_id(service_id):
            raise click.ClickException(
                "archive is not currently supported for v2 services."
                "Please contact Anyscale support for more information."
            )

        self.job_controller.archive(service_id, None, is_service=True)

    def rollback(self, service_id: str) -> None:
        if not _is_v2_service_id(service_id):
            raise click.ClickException(
                "rollback is only supported for v2 services."
                "Please contact Anyscale support for more information."
            )

        service = self.api_client.rollback_service_api_v2_services_v2_service_id_rollback_post(
            service_id
        ).result

        self.log.info(f"Service {service.id} rollback initiated.")
        self.log.info(
            f'View the service in the UI at {get_endpoint(f"/services/{service.id}")}.'
        )

    def terminate(self, service_id: str) -> None:
        if _is_v2_service_id(service_id):
            self.api_client.terminate_service_api_v2_services_v2_service_id_terminate_post(
                service_id
            )
        else:
            self.job_controller.terminate(service_id, None, is_service=True)
