from aws_cdk import core as cdk

import aws_cdk.aws_mwaa as mwaa

from aws_cdk import (
    core as cdk,
    aws_ec2 as ec2,
    aws_mwaa as mwaa,
    aws_s3 as s3,
    aws_s3_deployment as s3d,
    aws_iam as iam,
)


class MwaaStack(cdk.Stack):
    def __init__(self, scope: cdk.Construct, construct_id: str, **kwargs) -> None:
        super().__init__(scope, construct_id, **kwargs)

        # We'll create a VPC just for this
        vpc = ec2.Vpc(self, "mwaa-vpc", max_azs=2)

        # We need a bucket for assets
        bucket = s3.Bucket(
            self,
            "mwaa-bucket",
            versioned=True,
            auto_delete_objects=True,
            removal_policy=cdk.RemovalPolicy.DESTROY,
            block_public_access=s3.BlockPublicAccess.BLOCK_ALL,
        )
        files = s3d.BucketDeployment(
            self,
            "mwaa-assets",
            sources=[s3d.Source.asset("./assets")],
            destination_bucket=bucket,
        )

        # And a service role with additional EMR permissions
        # See https://docs.aws.amazon.com/mwaa/latest/userguide/mwaa-create-role.html
        mwaa_service_role = iam.Role(
            self,
            "mwaa-service-role",
            assumed_by=iam.CompositePrincipal(
                iam.ServicePrincipal("airflow.amazonaws.com"),
                iam.ServicePrincipal("airflow-env.amazonaws.com"),
            ),
            inline_policies={
                "CDKmwaaPolicyDocument": self.mwaa_policy_document(
                    "dacort-airflow", bucket.bucket_arn
                )
            },
            path="/service-role/",
        )

        # And security group
        security_group = ec2.SecurityGroup(
            self, id="mwaa-sg", vpc=vpc, security_group_name="mwaa-sg"
        )
        security_group.connections.allow_internally(ec2.Port.all_traffic(), "MWAA")

        # Enable logging on everything
        logging_configuration = mwaa.CfnEnvironment.LoggingConfigurationProperty(
            task_logs=mwaa.CfnEnvironment.ModuleLoggingConfigurationProperty(
                enabled=True, log_level="INFO"
            ),
            worker_logs=mwaa.CfnEnvironment.ModuleLoggingConfigurationProperty(
                enabled=True, log_level="INFO"
            ),
            scheduler_logs=mwaa.CfnEnvironment.ModuleLoggingConfigurationProperty(
                enabled=True, log_level="INFO"
            ),
            dag_processing_logs=mwaa.CfnEnvironment.ModuleLoggingConfigurationProperty(
                enabled=True, log_level="INFO"
            ),
            webserver_logs=mwaa.CfnEnvironment.ModuleLoggingConfigurationProperty(
                enabled=True, log_level="INFO"
            ),
        )

        # Create our MWAA
        subnets = [subnet.subnet_id for subnet in vpc.private_subnets]
        airflow = mwaa.CfnEnvironment(
            self,
            "airflow-v2",
            name="dacort-airflow",
            airflow_version="2.0.2",
            dag_s3_path=f"dags/",
            source_bucket_arn=bucket.bucket_arn,
            execution_role_arn=mwaa_service_role.role_arn,
            requirements_s3_path="requirements.txt",
            webserver_access_mode="PUBLIC_ONLY",
            environment_class="mw1.small",
            network_configuration=mwaa.CfnEnvironment.NetworkConfigurationProperty(
                subnet_ids=subnets,
                security_group_ids=[security_group.security_group_id],
            ),
            logging_configuration=logging_configuration,
        )
        airflow.node.add_dependency(files)

        # Register a couple outputs
        cdk.CfnOutput(self, "mwaa_bucket", value=bucket.bucket_name)
        cdk.CfnOutput(self, "mwaa_url", value=f"https://{airflow.attr_webserver_url}")

    def mwaa_policy_document(self, mwaa_env_name: str, mwaa_bucket_arn: str):
        return iam.PolicyDocument(
            statements=[
                iam.PolicyStatement(
                    actions=["airflow:PublishMetrics"],
                    effect=iam.Effect.ALLOW,
                    resources=[
                        f"arn:aws:airflow:{self.region}:{self.account}:environment/{mwaa_env_name}"
                    ],
                ),
                iam.PolicyStatement(
                    actions=["s3:ListAllMyBuckets"],
                    effect=iam.Effect.DENY,
                    resources=[f"{mwaa_bucket_arn}/*", f"{mwaa_bucket_arn}"],
                ),
                iam.PolicyStatement(
                    actions=["s3:*"],
                    effect=iam.Effect.ALLOW,
                    resources=[f"{mwaa_bucket_arn}/*", f"{mwaa_bucket_arn}"],
                ),
                iam.PolicyStatement(
                    actions=[
                        "logs:CreateLogStream",
                        "logs:CreateLogGroup",
                        "logs:PutLogEvents",
                        "logs:GetLogEvents",
                        "logs:GetLogRecord",
                        "logs:GetLogGroupFields",
                        "logs:GetQueryResults",
                    ],
                    effect=iam.Effect.ALLOW,
                    resources=[
                        f"arn:aws:logs:{self.region}:{self.account}:log-group:airflow-{mwaa_env_name}-*"
                    ],
                ),
                iam.PolicyStatement(
                    actions=["logs:DescribeLogGroups"],
                    effect=iam.Effect.ALLOW,
                    resources=["*"],
                ),
                iam.PolicyStatement(
                    actions=[
                        "sqs:ChangeMessageVisibility",
                        "sqs:DeleteMessage",
                        "sqs:GetQueueAttributes",
                        "sqs:GetQueueUrl",
                        "sqs:ReceiveMessage",
                        "sqs:SendMessage",
                    ],
                    effect=iam.Effect.ALLOW,
                    resources=[f"arn:aws:sqs:{self.region}:*:airflow-celery-*"],
                ),
                iam.PolicyStatement(
                    actions=[
                        "kms:Decrypt",
                        "kms:DescribeKey",
                        "kms:GenerateDataKey*",
                        "kms:Encrypt",
                    ],
                    effect=iam.Effect.ALLOW,
                    resources=["*"],
                    conditions={
                        "StringEquals": {
                            "kms:ViaService": [
                                f"sqs.{self.region}.amazonaws.com",
                                f"s3.{self.region}.amazonaws.com",
                            ]
                        }
                    },
                ),
                iam.PolicyStatement(
                    actions=[
                        "emr-containers:StartJobRun",
                        "emr-containers:DescribeJobRun",
                        "emr-containers:CancelJobRun",
                    ],
                    effect=iam.Effect.ALLOW,
                    resources=["*"],
                ),
                iam.PolicyStatement(
                    actions=[
                        "elasticmapreduce:RunJobFlow",
                        "elasticmapreduce:DescribeStep",
                        "elasticmapreduce:DescribeCluster",
                    ],
                    effect=iam.Effect.ALLOW,
                    resources=["*"],
                ),
                iam.PolicyStatement(
                    actions=["iam:PassRole"],
                    effect=iam.Effect.ALLOW,
                    resources=[
                        f"arn:aws:iam::{self.account}:role/EMR_DemoRole",
                        f"arn:aws:iam::{self.account}:role/EMR_EC2_DemoRole",
                        f"arn:aws:iam::{self.account}:role/EMR_EC2_DefaultRole",
                        f"arn:aws:iam::{self.account}:role/EMR_DefaultRole",
                    ],
                ),
            ]
        )
