import ipaddress
import re
from typing import Iterable, List, Tuple

from google.api_core.exceptions import NotFound, PermissionDenied
from google.cloud import compute_v1
from google.cloud.filestore_v1.types import Instance
from google.cloud.storage.bucket import Bucket
from google.cloud.storage.constants import (
    PUBLIC_ACCESS_PREVENTION_ENFORCED,
    PUBLIC_ACCESS_PREVENTION_INHERITED,
)
from google.iam.v1.policy_pb2 import Binding
from googleapiclient.errors import HttpError

from anyscale.cli_logger import BlockLogger
from anyscale.client.openapi_client.models import CreateCloudResourceGCP
from anyscale.shared_anyscale_utils.conf import ANYSCALE_HOST
from anyscale.utils.gcp_utils import (
    binding_from_dictionary,
    check_policy_bindings,
    check_required_policy_bindings,
    GCPLogger,
    GoogleCloudClientFactory,
)
from anyscale.utils.network_verification import (
    check_inbound_firewall_permissions,
    FirewallRule,
    GCP_SUBNET_CAPACITY,
    Protocol,
)


# Refer to https://cloud.google.com/python/docs/reference/file/latest/google.cloud.filestore_v1.
# services.cloud_filestore_manager.CloudFilestoreManagerClient, filestore name is in the format
# projects/{project_id}/locations/{location}/instances/{instance_id}.
_FILESTORE_NAME_REGEX_PATTERN = r"projects/(?P<project_id>[^/]+)/locations/(?P<location>[^/]+)/instances/(?P<instance_id>[^/]+)"


def verify_gcp_networking(
    factory: GoogleCloudClientFactory,
    resources: CreateCloudResourceGCP,
    project_id: str,
    cloud_region: str,
    logger: GCPLogger,
) -> bool:
    """Verify the existence and connectedness of the VPC & Subnet."""
    vpc_name = resources.gcp_vpc_id
    # TODO Verify Internet Gateway
    try:
        vpc = factory.compute_v1.NetworksClient().get(
            project=project_id, network=vpc_name
        )
    except NotFound:
        logger.log_resource_not_found_error("VPC", vpc_name)
        return False

    subnet_name = resources.gcp_subnet_ids[
        0
    ]  # TODO (congding): multiple subnets provided
    if len(resources.gcp_subnet_ids) > 1:
        logger.internal.warning(
            "Multiple subnets provided. Only taking the first subnet and ignoring the rest."
        )
    try:
        subnet = factory.compute_v1.SubnetworksClient().get(
            project=project_id, subnetwork=subnet_name, region=cloud_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, project_id: str, 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()
    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: CreateCloudResourceGCP,
    project_id: str,
    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.
    """
    anyscale_access_service_account = resources.gcp_anyscale_iam_service_account_email

    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 service_account_iam_policy.get("bindings") is None or 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."
        )

    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 True


def verify_gcp_dataplane_service_account(
    factory: GoogleCloudClientFactory,
    resources: CreateCloudResourceGCP,
    project_id: str,
    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.gcp_cluster_node_service_account_email
    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

    if resp["projectId"] != project_id:
        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"], project_id
            )
        )

    project_client = factory.resourcemanager_v3.ProjectsClient()
    project_iam_bindings = project_client.get_iam_policy(
        resource=f"projects/{project_id}"
    )
    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 {project_id}.\n"
            "This is safe to ignore if you are not using your own container images or are configuring access a different way."
        )

    return True


def verify_firewall_policy(
    factory: GoogleCloudClientFactory,
    resources: CreateCloudResourceGCP,
    project_id: str,
    cloud_region: str,
    logger: GCPLogger,
) -> bool:
    """Checks if the given firewall exists at either the Global or Regional level."""
    firewall_policy = resources.gcp_firewall_policy_ids[
        0
    ]  # TODO (congding): multiple firewall ids provided
    if len(resources.gcp_firewall_policy_ids) > 1:
        logger.internal.warning(
            "Multiple firewall policies provided. Only taking the first firewall policy and ignoring the rest."
        )
    vpc_name = resources.gcp_vpc_id

    firewall = compute_v1.types.compute.FirewallPolicy()
    try:
        firewall = factory.compute_v1.NetworkFirewallPoliciesClient().get(
            project=project_id, firewall_policy=firewall_policy
        )
    except NotFound:
        logger.internal.info(
            f"Global firweall policy {firewall_policy} not found, trying in region: {cloud_region}."
        )
    if not firewall:
        try:
            firewall = factory.compute_v1.RegionNetworkFirewallPoliciesClient().get(
                project=project_id, firewall_policy=firewall_policy, region=cloud_region
            )
        except NotFound:
            logger.log_resource_not_found_error("FirewallPolicy", firewall_policy)
            return False

    if not any(
        association.attachment_target.split("/")[-1] == vpc_name
        for association in firewall.associations
    ):
        logger.internal.error(
            "Firewall policy {} is not associated with the VPC {}, but is associated with the following VPCs: {}".format(
                firewall_policy, vpc_name, firewall.associations
            )
        )
        return False

    subnet_obj = factory.compute_v1.SubnetworksClient().get(
        project=project_id, subnetwork=resources.gcp_subnet_ids[0], region=cloud_region
    )  # TODO (congding): multiple subnets provided
    subnet = ipaddress.ip_network(subnet_obj.ip_cidr_range)

    rules = _firewall_rules_from_proto_resp(firewall.rules)
    if not check_inbound_firewall_permissions(
        rules, Protocol.tcp, {22}, ipaddress.ip_network("0.0.0.0/0")
    ):
        logger.internal.warning(
            "Firewall policy does not allow inbound access on port 22. Certain features may be unavailable."
        )
    if not check_inbound_firewall_permissions(
        rules, Protocol.tcp, {443}, ipaddress.ip_network("0.0.0.0/0")
    ):
        logger.internal.warning(
            "Firewall policy does not allow inbound access on port 443. Accessing features like Jupyter & Ray Dashboard may not work."
        )

    if not check_inbound_firewall_permissions(rules, Protocol.tcp, None, subnet):
        logger.internal.error(
            "Firewall policy does not allow for internal communication, see https://cloud.google.com/vpc/docs/using-firewalls#common-use-cases-allow-internal for how to configure such a rule."
        )
        return False
    return True


def _firewall_rules_from_proto_resp(
    rules: Iterable[compute_v1.types.compute.FirewallPolicyRule],
) -> List[FirewallRule]:
    """Convert Firewall Rules into a Python class suitable for comparison.

    NOTE: We only check "ALLOW" rules & do not check priority.
    """
    return [
        FirewallRule(
            direction=rule.direction,
            protocol=Protocol.from_val(l4config.ip_protocol),
            ports=l4config.ports,
            network=ipaddress.ip_network(addr),
        )
        for rule in rules
        if rule.action == "allow"
        for addr in rule.match.src_ip_ranges
        for l4config in (
            rule.match.layer4_configs
            or [
                compute_v1.types.compute.FirewallPolicyRuleMatcherLayer4Config(
                    ip_protocol="all"
                )
            ]
        )
    ]


def verify_filestore(
    factory: GoogleCloudClientFactory,
    resources: CreateCloudResourceGCP,
    cloud_region: str,
    logger: GCPLogger,
) -> bool:
    """Verify Filestore exists & that it is connected to the correct VPC.

    TODO: Warn about Filestore size if it is 'too small'."""
    file_store_instance_name = resources.gcp_filestore_config.instance_name
    client = factory.filestore_v1.CloudFilestoreManagerClient()
    try:
        file_store = client.get_instance(name=file_store_instance_name)
    except NotFound:
        logger.log_resource_not_found_error("Filestore", file_store_instance_name)
        return False
    # file_store 1.6.0: https://cloud.google.com/python/docs/reference/file/latest/google.cloud.filestore_v1.types.Instance.Tier
    if file_store.tier != Instance.Tier.ENTERPRISE:
        logger.internal.warning(
            f"Filestore is running with tier {file_store.tier.name}, {Instance.Tier.ENTERPRISE.name} is the suggested tier."
        )

    file_store_match = re.match(_FILESTORE_NAME_REGEX_PATTERN, file_store.name)
    if file_store_match:
        filestore_location = file_store_match.group("location")
        # Filestore location can be region or zone.
        if cloud_region not in filestore_location:
            logger.internal.warning(
                f"Filestore is in {filestore_location}, but this cloud is being set up in {cloud_region}. This can result in cross-region network egress charges, refer to https://cloud.google.com/vpc/network-pricing#egress-within-gcp."
            )
    else:
        logger.internal.error(
            f"Failed to get filestore location, since filestore name doesn't match format {_FILESTORE_NAME_REGEX_PATTERN}!"
        )

    file_store_networks = [v.network for v in file_store.networks]
    vpc_name = resources.gcp_vpc_id
    if vpc_name not in file_store_networks:
        logger.internal.error(
            f"Filestore is not connected to {vpc_name}, but to {file_store_networks}. "
            f"This cannot be edited on an existing Filestore instance. Please recreate the filestore and connect it to {vpc_name}."
        )
        return False

    return True


def verify_cloud_storage(
    factory: GoogleCloudClientFactory,
    resources: CreateCloudResourceGCP,
    project_id: str,
    cloud_region: str,
    logger: GCPLogger,
):
    """Verify that the Google Cloud Storage Bucket exists & raises warnings about improper configurations."""
    bucket_client = factory.storage.Client(project_id)
    bucket_name = resources.gcp_cloud_storage_bucket_id
    try:
        bucket = bucket_client.get_bucket(bucket_name)
    except NotFound:
        logger.log_resource_not_found_error("Google Cloud Storage Bucket", bucket_name)
        return False

    if not bucket.iam_configuration.uniform_bucket_level_access_enabled:
        logger.internal.warning(
            f"Bucket {bucket_name}, does not have Uniform Bucket Access enabled, "
            "this impedes Anyscale's ability to verify bucket access."
        )

    public_access_protection = bucket.iam_configuration.public_access_prevention
    if public_access_protection not in (
        PUBLIC_ACCESS_PREVENTION_INHERITED,
        PUBLIC_ACCESS_PREVENTION_ENFORCED,
    ):
        logger.internal.warning(
            f"Bucket {bucket_name} has public access prevention set to {public_access_protection}, not enforced or inherited."
        )

    correct_region, bucket_region = _check_bucket_region(bucket, cloud_region)
    if not correct_region:
        logger.internal.warning(
            f"Bucket {bucket_name} is in region {bucket_region}, but this cloud is being set up in {cloud_region}."
            "This can result in degraded cluster launch & logging performance."
        )

    if not any(
        (
            ANYSCALE_HOST in cors_config.get("origin")
            and "*" in cors_config.get("responseHeader", [])
            and "GET" in cors_config.get("method", [])
        )
        for cors_config in bucket.cors
    ):
        logger.internal.warning(
            f"Bucket {bucket_name} does not have the correct CORS rule for Anyscale. This is safe to ignore if you are not using Anyscale UI.\n"
            "If you are using the UI, please create the correct CORS rule for Anyscale according to https://docs.anyscale.com/user-guide/onboard/clouds/deploy-on-gcp#appendix-detailed-resource-requirements"
        )

    iam_policy = bucket.get_iam_policy()
    iam_bindings = binding_from_dictionary(iam_policy.bindings)

    permission_warning = (
        "The {location} Service Account {email} requires the following permissions on Bucket {bucket} to operate correctly:\n"
        "* storage.buckets.get\n* storage.objects.[ get | list | create ]\n* storage.multipartUploads.[ abort | create | listParts ]\n"
        "We suggest granting the predefined roles of `roles/storage.legacyBucketReader` and `roles/storage.objectAdmin`."
    )

    access_service_account = resources.gcp_anyscale_iam_service_account_email
    if not _verify_service_account_on_bucket(
        f"serviceAccount:{access_service_account}", iam_bindings
    ):
        logger.confirm_missing_permission(
            permission_warning.format(
                location="Anyscale access",
                email=access_service_account,
                bucket=bucket_name,
            )
        )

    dataplane_service_account = resources.gcp_cluster_node_service_account_email
    if not _verify_service_account_on_bucket(
        f"serviceAccount:{dataplane_service_account}", iam_bindings
    ):
        logger.confirm_missing_permission(
            permission_warning.format(
                location="Dataplane",
                email=dataplane_service_account,
                bucket=bucket_name,
            )
        )

    return True


def _verify_service_account_on_bucket(
    service_account: str, iam_bindings: List[Binding]
) -> bool:
    """Verifies that the given service account has roles that ensure the following list of permissions:
    * storage.buckets.get: Get bucket info
    * storage.objects.[ get | list | create ]
    * storage.multipartUploads.[ abort | create | listParts ]
    """
    if check_policy_bindings(iam_bindings, service_account, {"roles/storage.admin"},):
        return True
    return any(
        check_required_policy_bindings(iam_bindings, service_account, combination)
        for combination in [
            {
                # Best Permissions (legacy for `buckets.get`)
                "roles/storage.legacyBucketReader",
                "roles/storage.objectViewer",
                "roles/storage.objectCreator",
            },
            {
                # Most Brief permissions (legacy for `buckets.get`)
                "roles/storage.objectAdmin",
                "roles/storage.legacyBucketReader",
            },
            {
                # Optimal Legacy Roles
                "roles/storage.legacyBucketWriter",
                "roles/storage.legacyObjectReader",
            },
            {
                # Legacy Roles (extra object permissions)
                "roles/storage.legacyBucketWriter",
                "roles/storage.legacyObjectOwner",
            },
            {
                # Legacy Roles (extra bucket permissions)
                "roles/storage.legacyBucketOwner",
                "roles/storage.legacyObjectReader",
            },
            {
                # Legacy Roles (extra bucket/object permissions)
                "roles/storage.legacyBucketOwner",
                "roles/storage.legacyObjectOwner",
            },
        ]
    )


def _check_bucket_region(bucket: Bucket, region: str) -> Tuple[bool, str]:
    if bucket.location_type == "dual-region":
        return region.upper() in bucket.data_locations, ",".join(bucket.data_locations)
    elif bucket.location_type == "region":
        return region.upper() == bucket.location, bucket.location

    # Bucket is `multi-region`, so check if the location (`EU`, `ASIA`, `US`)
    # is the 'prefix' of the region.
    return region.upper().startswith(bucket.location), bucket.location
