import aws_cdk as cdk
from constructs import Construct
from aws_cdk import (
    aws_s3 as _s3,
    aws_iam as _iam,
    aws_cognito as _cognito,
    aws_kms as _kms,
)
import boto3
from aws_cdk_constructs.utils import normalize_environment_parameter, get_version
from cloudcomponents.cdk_cloudfront_authorization import StaticSiteAuthorization, StaticSiteDistribution
import aws_cdk.aws_cloudfront_origins as origins


def check_if_bucket_exist(s3_bucket_name, profile=None):
    aws_session = boto3.Session(profile_name=profile)
    sdk_s3 = aws_session.client("s3")
    bucket_exist = True

    if not s3_bucket_name:
        return False

    try:
        sdk_s3.get_bucket_location(Bucket=s3_bucket_name)
    except:
        bucket_exist = False

    return bucket_exist


class Bucket(Construct):
    """ Construct to create an S3 Bucket

    Args:

        id (str): the logical id of the newly created resource

        app_name (str): The application name. This will be used to generate the 'ApplicationName' tag for CSI compliancy. The ID of the application. This must be unique for each system, as it will be used to calculate the AWS costs of the system

        environment (str): Specify the environment in which you want to deploy you system. Allowed values: Development, QA, Production, SharedServices 

        environments_parameters (dict): The dictionary containing the references to CSI AWS environments. This will simplify the environment promotions and enable a parametric development of the infrastructures.

        bucket_name (str): The S3 Bucket in which the application S3 are stored. In case the bucket is configure to have CDN integrated with Cognito, this is also the CloudFront Distribution alias (i.e. the final domain of the website)

        bucket_is_public (str): Wheather or not the S3 bucket should be public

        bucket_is_privately_accessed_from_vpc_over_http_over_http (str): Force the bucket to be private and enable HTTP private accessed from within the VPC. When this parameter is set to True, `bucket_is_public` and `bucket_has_cdn` will be forced to be False

        bucket_has_cdn (str): Wheather or not the S3 bucket will be serverd by a Cloudflare CDN

        bucket_logic_id (str): The logical id of the S3 Bucket

        bucket_website_index_document (str): Use this parameter to configure the S3 bucket as Web Hosting. This is the S3 key of the index document of your static site (generally is index.html)

        bucket_website_error_document (str): Use this parameter to configure the S3 bucket as Web Hosting. This is the S3 key of the error document of your static site (generally is error.html)

        bucket_has_authentication (str): Create a CloudFront distribuition integrated with Cognito, to enable authenticated access to the bucket
        
        bucket_is_encrypted (str): Whether this bucket should be encrypted. KMS is autogenerated. Default: None
        
        versioned (str): Whether this bucket should have versioning turned on or not. Default: false

        encryption_allowed_principals_arns (list | str): A single ARN or an array of ARNs to be added as principal to the KMS encryption key policy of the bucket, when encrypted. These principals can utilize the key to interact with the encrypted bucket.

        cdn_http_headers ([string: string]): A dictionary of HTTP headers to be added to the CloudFront distribution.
    """
    @property
    def get_s3_bucket(self):
        """Returns the S3 bucket

        Returns:
            aws_s3.Bucket: the S3 bucket
        """
        return self.bucket

    def create_bucket(
        self,
        logic_id,
        bucket_name,
        app_name,
        access_control=None,
        public_read_access=False,
        website_index_document=None,
        website_error_document=None,
        removal_policy=None,
        versioned=False,
        is_encrypted=False,
        encryption_allowed_principals_arns=None,
        **kwargs
    ):
        """Create an S3 bucket

        Args:

            logic_id (str): The logical ID of the S3 Bucket

            bucket_name (str): The S3 Bucket name

            app_name (str): The application name. This will be used to generate the 'ApplicationName' tag for CSI compliancy. The ID of the application. This must be unique for each system, as it will be used to calculate the AWS costs of the system

            access_control (aws_cdk.aws_s3.BucketAccessControl): The S3 Bucket access control policy. For more info see `access_control` in https://docs.aws.amazon.com/cdk/api/latest/python/aws_cdk.aws_s3/Bucket.html

            public_read_access (bool): Whether the S3 Bucket should be public or not

            website_index_document (str): Use this parameter to configure the S3 bucket as Web Hosting. This is the S3 key of the index document of your static site (generally is index.html)

            website_error_document (str): Use this parameter to configure the S3 bucket as Web Hosting. This is the S3 key of the error document of your static site (generally is error.html)

            removal_policy (aws_cdk. cdk.RemovalPolicy): The S3 Bucket removal policy. For more info see `removal_policy` in https://docs.aws.amazon.com/cdk/api/latest/python/aws_cdk.aws_s3/Bucket.html

            versioned (Bool): Whether this bucket should have versioning turned on or not. Default: false

            is_encrypted (Bool): Whether this bucket should be encrypted. KMS is autogenerated. Default: false

            encryption_allowed_principals_arns (list | str): A single ARN or an array of ARNs to be added as principal to the KMS encryption key policy of the bucket, when encrypted. These principals can utilize the key to interact with the encrypted bucket.

        Returns:
            aws_s3.Bucket: the S3 bucket
        """
        # Create S3 bucket
        s3_bucket = _s3.Bucket(
            self,
            logic_id,
            bucket_name=bucket_name,
            access_control=access_control,
            public_read_access=public_read_access,
            website_index_document=website_index_document,
            website_error_document=website_error_document,
            removal_policy=removal_policy,
            versioned=versioned,
            encryption=_s3.BucketEncryption.KMS if is_encrypted else None
        )

        if is_encrypted:
            key = s3_bucket.node.find_child("Key")
            _kms.Alias(
                self, 
                logic_id + "kmskeyalias",
                alias_name=app_name + "-" + logic_id,
                target_key=key)

            if encryption_allowed_principals_arns:
                encryption_bucket_policy = _iam.PolicyStatement(
                    actions=[
                        "kms:Create*",
                        "kms:Describe*",
                        "kms:Decrypt",
                        "kms:Enable*",
                        "kms:List*",
                        "kms:Put*",
                        "kms:Update*",
                        "kms:Revoke*",
                        "kms:Disable*",
                        "kms:Get*",
                        "kms:Delete*",
                        "kms:ScheduleKeyDeletion",
                        "kms:CancelKeyDeletion",
                        "kms:GenerateDataKey",
                        "kms:TagResource",
                        "kms:UntagResource"
                    ],
                    resources=["*"],
                )
                
                # Check if `encryption_allowed_principals_arns` is an array, otherwise convert it
                if not isinstance(encryption_allowed_principals_arns, list):
                    encryption_allowed_principals_arns = [encryption_allowed_principals_arns]

                # Loop over the arns arrays and add them all to the statement
                for arns in encryption_allowed_principals_arns:
                    encryption_bucket_policy.add_arn_principal(arns)
                
                key.add_to_resource_policy(statement=encryption_bucket_policy)   
  
        cdk.Tags.of(s3_bucket).add("ApplicationName", app_name,)

        return s3_bucket

    def __init__(
        self,
        scope: Construct,
        id: str,
        app_name,
        environment,
        environments_parameters,
        bucket_name,
        bucket_logic_id=None,
        bucket_is_public=None,
        bucket_has_cdn=None,
        bucket_website_index_document=None,
        bucket_website_error_document=None,
        bucket_is_privately_accessed_from_vpc_over_http=False,
        bucket_has_authentication=None,
        bucket_is_encrypted=None,
        versioned=False,
        encryption_allowed_principals_arns=None,
        cdn_http_headers=None,
        **kwargs
    ):
        super().__init__(scope, id, **kwargs)
        environment = normalize_environment_parameter(environment)
        app_name = app_name.lower().strip()

        # Apply mandatory tags
        cdk.Tags.of(self).add("ApplicationName", app_name)
        cdk.Tags.of(self).add("Environment", environment)

        # Apply FAO CDK tags
        cdk.Tags.of(self).add("fao-cdk-construct", "bucket")
        cdk.Tags.of(cdk.Stack.of(self)).add("fao-cdk-version", get_version())
        cdk.Tags.of(cdk.Stack.of(self)).add("fao-cdk", "true")

        # Declare variables
        self.bucket = None

        environment = environment.lower()
        aws_account = environments_parameters["accounts"][environment]
        account_id = aws_account["id"]
        logic_id = bucket_logic_id if bucket_logic_id else "S3"

        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Create conditions
        # bucket_name_was_provided = bucket_name.strip()
        # bucket_already_exist = check_if_bucket_exist(
        #     s3_bucket_name=bucket_name
        # )

        # Check if bucket has to be public
        public_read_access = (
            bucket_is_public
            and isinstance(bucket_is_public, str)
            and bucket_is_public.lower() == "true"
        )

        include_cdn = (
            bucket_has_cdn
            and isinstance(bucket_has_cdn, str)
            and bucket_has_cdn.lower() == "true"
        )

        has_authentication = (
            bucket_has_authentication
            and isinstance(bucket_has_authentication, str)
            and bucket_has_authentication.lower() == "true"
        )

        is_production = environment.lower().strip() == "production"
        is_not_production = not is_production

        is_encrypted = (
            bucket_is_encrypted
            and isinstance(bucket_is_encrypted, str)
            and bucket_is_encrypted.lower() == "true"
        )

        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ CloudFormation outputs
        cdk.CfnOutput(self, f"BucketAppName{app_name}{bucket_name}", value=str(app_name))
        cdk.CfnOutput(self, f"BucketPublicReadAccess{app_name}{bucket_name}", value=str(public_read_access))
        cdk.CfnOutput(self, f"BucketIncludesCdn{app_name}{bucket_name}", value=str(include_cdn))
        cdk.CfnOutput(self, f"BucketHasAuthentication{app_name}{bucket_name}", value=str(has_authentication))
        cdk.CfnOutput(self, f"BucketIsEncrypted{app_name}{bucket_name}", value=str(is_encrypted))
        cdk.CfnOutput(self, f"BucketBucketName{app_name}{bucket_name}", value=str(bucket_name))


        # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Conditionally create resources

        # Only retain production bucket, with public assets
        removal_policy =  cdk.RemovalPolicy.DESTROY
        if is_production and public_read_access:
            removal_policy =  cdk.RemovalPolicy.RETAIN

        # Default access level: PRIVATE
        self.bucket = self.create_bucket(
            logic_id=logic_id,
            app_name=app_name,
            bucket_name=bucket_name,
            website_index_document=bucket_website_index_document,
            website_error_document=bucket_website_error_document,
            removal_policy=removal_policy,
            access_control=_s3.BucketAccessControl.PRIVATE if not public_read_access else None,
            # https://docs.aws.amazon.com/AmazonS3/latest/dev/access-control-block-public-access.html,
            block_public_access=_s3.BlockPublicAccess.BLOCK_ALL if not public_read_access else None,
            versioned=versioned,
            is_encrypted=is_encrypted,
            encryption_allowed_principals_arns=encryption_allowed_principals_arns
        )

        if bucket_is_privately_accessed_from_vpc_over_http:
            # Force public read to false
            public_read_access = False
            bucket_has_cdn = False

            self.bucket.add_to_resource_policy(
                _iam.PolicyStatement(
                    principals=[_iam.AnyPrincipal()],
                    actions=["s3:GetObject"],
                    resources=[
                        "arn:aws:s3:::" + bucket_name,
                        "arn:aws:s3:::" + bucket_name + "/*",
                    ],
                    #conditions={"StringEquals": {"aws:sourceVpce": aws_account["vpc"]}},
                )
            )

        if public_read_access:
            # Grant access to anyone
            self.bucket.grant_public_access()

            self.bucket.add_cors_rule(
                allowed_methods=[
                    _s3.HttpMethods.GET,
                    _s3.HttpMethods.POST,
                    _s3.HttpMethods.PUT,
                    _s3.HttpMethods.DELETE,
                    _s3.HttpMethods.HEAD,
                ],
                allowed_origins=["*"],
                allowed_headers=["*"],
            )

        # Limit access to CDN
        if include_cdn:

            # Cloudflare doc https://support.cloudflare.com/hc/en-us/articles/360037983412-Configuring-an-Amazon-Web-Services-static-site-to-use-Cloudflare
            cloudflare_ips = [
                "2400:cb00::/32",
                "2405:8100::/32",
                "2405:b500::/32",
                "2606:4700::/32",
                "2803:f800::/32",
                "2c0f:f248::/32",
                "2a06:98c0::/29",
                "103.21.244.0/22",
                "103.22.200.0/22",
                "103.31.4.0/22",
                "104.16.0.0/13",
                "104.24.0.0/14",
                "108.162.192.0/18",
                "131.0.72.0/22",
                "141.101.64.0/18",
                "162.158.0.0/15",
                "172.64.0.0/13",
                "173.245.48.0/20",
                "188.114.96.0/20",
                "190.93.240.0/20",
                "197.234.240.0/22",
                "198.41.128.0/17",
            ]

            # Only allow access from Cloudflare IPs
            self.bucket.add_to_resource_policy(
                _iam.PolicyStatement(
                    effect=_iam.Effect.DENY,
                    principals=[_iam.AnyPrincipal()],
                    actions=["s3:GetObject"],
                    resources=[
                        "arn:aws:s3:::" + bucket_name + "/*",
                    ],
                    conditions={"NotIpAddress": {
                        "aws:SourceIp": cloudflare_ips}},
                )
            )
            self.bucket.add_to_resource_policy(
                _iam.PolicyStatement(
                    effect=_iam.Effect.ALLOW,
                    principals=[_iam.AnyPrincipal()],
                    actions=["s3:GetObject"],
                    resources=[
                        "arn:aws:s3:::" + bucket_name + "/*",
                    ],
                    conditions={"IpAddress": {"aws:SourceIp": cloudflare_ips}},
                )
            )

            self.bucket.add_cors_rule(
                allowed_methods=[
                    _s3.HttpMethods.GET,
                    _s3.HttpMethods.POST,
                    _s3.HttpMethods.PUT,
                    _s3.HttpMethods.DELETE,
                    _s3.HttpMethods.HEAD,
                ],
                allowed_origins=["*"],
                allowed_headers=["*"],
            )

        if has_authentication:

            # Dynamically retrieve the user pool according to FAO environment
            user_pool_arn = aws_account["cognito_user_pool_arn"]

            user_pool = _cognito.UserPool.from_user_pool_arn(
                self, "UserPool", user_pool_arn=user_pool_arn)

            authorization = StaticSiteAuthorization(self, "Authorization",
                user_pool=user_pool,
                identity_providers=[_cognito.UserPoolClientIdentityProvider.custom("FAO-SSO")],
                http_headers= cdn_http_headers if cdn_http_headers else None
            )
        
            static_distribution = StaticSiteDistribution(self, "Distribution",
                authorization=authorization,
                origin=origins.S3Origin(self.bucket),
                domain_names=[bucket_name],
            )

            # Set CloudFront Distribution alias and SSL certificate
            distribution_config = static_distribution.node.find_child("Distribution").node.find_child("Resource")

            distribution_config.add_override("Properties.DistributionConfig.Aliases", [bucket_name])
            distribution_config.add_override("Properties.DistributionConfig.ViewerCertificate.AcmCertificateArn", aws_account["ssl_certificate_star_fao_org_arn"])
            distribution_config.add_override("Properties.DistributionConfig.ViewerCertificate.SslSupportMethod", "sni-only")
