from functools import partial
from typing import Any, Dict

from google.api_core.exceptions import NotFound, PermissionDenied
from google.auth.credentials import Credentials
import google.cloud
from google.cloud import compute_v1, resourcemanager_v3  # noqa
from googleapiclient.discovery import build as api_client_build
from googleapiclient.errors import HttpError

from anyscale.cli_logger import BlockLogger
from anyscale.util import confirm
from anyscale.utils.gcp_utils import (
    binding_from_dictionary,
    check_policy_bindings,
    get_application_default_credentials,
)
from anyscale.utils.network_verification import GCP_SUBNET_CAPACITY


class GCPLogger:
    def __init__(self, logger: BlockLogger, project_id: str, yes: bool = False):
        self.internal = logger
        self.project_id = project_id
        self.yes = yes

    def log_resource_not_found_error(self, resource_name: str, resource_id: str):
        self.internal.error(
            f"Could not find {resource_name} with id {resource_id} in project {self.project_id}. Please validate that you're using the correct GCP project and that the resource values are correct."
        )

    def confirm_missing_permission(self, error_str: str):
        self.internal.error(error_str)
        confirm(
            "If the serivce account has these required permissions granted via other\n"
            "roles, press 'y' to continue with verification or 'N' to abort.",
            yes=self.yes,
        )


class GoogleCloudClientFactory:
    """Factory to generate both Google Cloud Client libraries & Google API Client libraries.

    Google Cloud Client libraries are instantiated by:
    ```
        factory = GoogleCloudClientFactory(credentials=AnonymousCredentials())
        client = factory.compute_v1.ExampleClient()
    ```

    Google API Client libraries are instantiated by:
    ```
        factory = GoogleCloudClientFactory(credentials=AnonymousCredentials())
        client = factory.build("iam", "v1")
    ```
    """

    def __init__(self, credentials: Credentials, force_rest=False, **kwargs):
        kwargs["credentials"] = credentials
        self.kwargs = kwargs
        self.force_rest = force_rest

    def __getattr__(self, client_library: str):
        """Get a wrapped Google Cloud Client library that injects default values from the factory."""
        module = getattr(google.cloud, client_library)
        kwargs = self.kwargs
        if self.force_rest:
            kwargs["transport"] = "rest"

        class WrappedClient:
            def __getattr__(self, client_type: str):
                return partial(getattr(module, client_type), **kwargs)

        return WrappedClient()

    def build(self, service_name: str, version: str):
        """Return a Google API Client with default values from the factor"""
        return api_client_build(
            service_name, version, cache_discovery=False, **self.kwargs
        )


def verify_gcp(resources: Dict[str, Any], logger: BlockLogger) -> bool:
    credentials, credentials_project = get_application_default_credentials(logger)
    specified_project = resources["project_id"]
    if credentials_project != specified_project:
        logger.warning(
            f"Default credentials are for {credentials_project}, but this cloud is being configured for {specified_project}"
        )

    factory = GoogleCloudClientFactory(credentials=credentials)
    gcp_logger = GCPLogger(logger, specified_project)

    return all(
        [
            _verify_gcp_project(factory, resources, gcp_logger),
            _verify_gcp_access_service_account(factory, resources, gcp_logger),
            _verify_gcp_dataplane_service_account(factory, resources, gcp_logger),
            _verify_gcp_networking(factory, resources, gcp_logger),
        ]
    )


def _verify_gcp_networking(
    factory: GoogleCloudClientFactory, resources: Dict, logger: GCPLogger,
) -> bool:
    """Verify the existence and connectedness of the VPC & Subnet."""

    project = resources["project_id"]
    vpc_name = resources["vpc_name"]
    # TODO Verify Internet Gateway
    try:
        vpc = factory.compute_v1.NetworksClient().get(project=project, network=vpc_name)
    except NotFound:
        logger.log_resource_not_found_error("VPC", vpc_name)
        return False

    subnet_name = resources["subnet_name"]
    try:
        subnet = factory.compute_v1.SubnetworksClient().get(
            project=project, subnetwork=subnet_name, region=resources["region"]
        )
    except NotFound:
        logger.log_resource_not_found_error("Subnet", subnet_name)
        return False

    if subnet.network != vpc.self_link:
        logger.internal.error(f"Subnet {subnet_name} is not part of {vpc_name}!")
        return False

    return _gcp_subnet_has_enough_capacity(subnet, logger.internal)


def _gcp_subnet_has_enough_capacity(
    subnet: compute_v1.types.compute.Subnetwork, logger: BlockLogger
) -> bool:
    """Verify if the subnet provided has a large enough IP address block."""
    if GCP_SUBNET_CAPACITY.verify_network_capacity(
        cidr_block_str=subnet.ip_cidr_range, resource_name=subnet.name, logger=logger,
    ):
        logger.info(f"Subnet {subnet.name} verification succeeded.")
        return True
    return False


def _verify_gcp_project(
    factory: GoogleCloudClientFactory, resources, logger: GCPLogger
) -> bool:
    """Verify if the project exists and is active.

    NOTE: This also checks that Compute Engine Service Agent is configured in the default way.
    """
    project_client = factory.resourcemanager_v3.ProjectsClient()
    project_id = resources["project_id"]
    try:
        project = project_client.get_project(name=f"projects/{project_id}")
    except (NotFound, PermissionDenied):
        logger.log_resource_not_found_error("Project", project_id)
        return False
    if project.state != project.State.ACTIVE:
        logger.internal.error(
            f"Project {project_id} is in state: {project.state}, not active"
        )
        return False

    iam_policies = project_client.get_iam_policy(resource=f"projects/{project_id}")
    project_number = project.name.split("/")[1]

    if not check_policy_bindings(
        iam_policies.bindings,
        f"serviceAccount:service-{project_number}@compute-system.iam.gserviceaccount.com",
        {"roles/compute.serviceAgent"},
    ):
        logger.internal.warning(
            "The Compute Engine Service Agent does not have the standard IAM Role of 'roles/compute.serviceAgent' in your project.\n"
            "This is not recommended by Google and may result in an inability for Compute Engine to function properly:\n"
            "https://cloud.google.com/compute/docs/access/service-accounts#compute_engine_service_account"
        )

    # TODO: Verify that APIs are Enabled. This will implicitly be verified during other verification steps.
    return True


def _verify_gcp_access_service_account(
    factory: GoogleCloudClientFactory, resources, logger: GCPLogger
) -> bool:
    """Verify if the Service Account meant to grant Anyscale access to this cloud has access to the specified project.

    NOTE: We verify that this service account can call signBlob on itself because this is necessary for downloading logs.
    """
    project_id = resources["project_id"]
    anyscale_access_service_account = resources["anyscale_access_service_account"]

    service_account_client = factory.build("iam", "v1").projects().serviceAccounts()
    try:
        service_account_iam_policy = service_account_client.getIamPolicy(
            resource=f"projects/-/serviceAccounts/{anyscale_access_service_account}"
        ).execute()
    except HttpError as e:
        if e.status_code == 404:
            logger.log_resource_not_found_error(
                "Anyscale Access Service Account", anyscale_access_service_account
            )
            return False

    if not check_policy_bindings(
        binding_from_dictionary(service_account_iam_policy["bindings"]),
        f"serviceAccount:{anyscale_access_service_account}",
        {"roles/iam.serviceAccountTokenCreator"},
    ):
        logger.confirm_missing_permission(
            f"Service Account {anyscale_access_service_account} must have the `iam.serviceAccounts.signBlob` permission to perform log download from the Anyscale platform.\n"
            "Please grant the `roles/iam.serviceAccountTokenCreator` role to this Service Account to pass this check."
        )
        return False

    project_client = factory.resourcemanager_v3.ProjectsClient()
    iam_policies = project_client.get_iam_policy(resource=f"projects/{project_id}")
    if not check_policy_bindings(
        iam_policies.bindings,
        f"serviceAccount:{anyscale_access_service_account}",
        {"roles/editor", "roles/owner"},
    ):
        logger.confirm_missing_permission(
            f"Service Account {anyscale_access_service_account} must have either `roles/editor` or `roles/owner` roles on the Project {project_id}.\n"
            "A more fine-grained set of permissions may be possible, but may break Anyscale's functionality."
        )
        return False

    return True


def _verify_gcp_dataplane_service_account(
    factory: GoogleCloudClientFactory, resources, logger: GCPLogger
) -> bool:
    """Verify that the Service Account for compute instances exists *in* the specified project.

    Compute Engine's ability to use this role
    This relies on the fact that Compute Engine Service Agent has the roles/compute.serviceAgent Role
    """
    service_account = resources["dataplane_service_account"]
    service_account_client = factory.build("iam", "v1").projects().serviceAccounts()
    try:
        resp = service_account_client.get(
            name=f"projects/-/serviceAccounts/{service_account}"
        ).execute()
    except HttpError as e:
        if e.status_code == 404:
            logger.log_resource_not_found_error(
                "Dataplane Service Account", service_account
            )
            return False
        raise e

    cloud_project = resources["project_id"]
    if resp["projectId"] != cloud_project:
        logger.internal.warning(
            "Service Account {} is in project {}, not {}. Please ensure that `constraints/iam.disableCrossProjectServiceAccountUsage` is not enforced.\n"
            "See: https://cloud.google.com/iam/docs/attach-service-accounts#enabling-cross-project for more information.".format(
                service_account, resp["projectId"], cloud_project
            )
        )

    project_client = factory.resourcemanager_v3.ProjectsClient()
    project_iam_bindings = project_client.get_iam_policy(
        resource=f"projects/{cloud_project}"
    )
    if not check_policy_bindings(
        project_iam_bindings.bindings,
        f"serviceAccount:{service_account}",
        {"roles/artifactregistry.reader", "roles/viewer"},
    ):
        logger.internal.warning(
            f"The dataplane Service Account {service_account} does not have `roles/artifactregistry.reader` or `roles/viewer` on {cloud_project}.\n"
            "This is safe to ignore if you are not using your own container images or are configuring access a different way."
        )

    return True
