# s3 bucker for mwaa dags
resource "aws_s3_bucket" "mwaa-scripts" {
  bucket = "${var.project_prefix}-mwaa-scripts-${var.env}"
  acl    = "private"
}

resource "aws_s3_bucket_public_access_block" "mwaa-scripts-access-block" {
  bucket = aws_s3_bucket.mwaa-scripts.id

  block_public_acls       = true
  block_public_policy     = true
  ignore_public_acls      = true
  restrict_public_buckets = true
}

# vpc for mwaa
resource "aws_vpc" "mwaa-vpc" {
  cidr_block = var.vpc_cidr
  tags = {
    Name = var.mwaa_env_name
  }
}

# vpc subnet for mwaa
resource "aws_subnet" "mwaa-private-1" {
  vpc_id            = aws_vpc.mwaa-vpc.id
  cidr_block        = var.private_subnet1_cidr
  availability_zone = var.region_az1
  tags = {
    Name = "${var.mwaa_env_name} Private Subnet 1"
  }
}

resource "aws_subnet" "mwaa-private-2" {
  vpc_id            = aws_vpc.mwaa-vpc.id
  cidr_block        = var.private_subnet2_cidr
  availability_zone = var.region_az2
  tags = {
    Name = "${var.mwaa_env_name} Private Subnet 2"
  }
}

resource "aws_subnet" "mwaa-public-1" {
  vpc_id            = aws_vpc.mwaa-vpc.id
  cidr_block        = var.public_subnet1_cidr
  availability_zone = var.region_az1
  map_public_ip_on_launch = true
  tags = {
    Name = "${var.mwaa_env_name} Public Subnet 1"
  }
}

resource "aws_subnet" "mwaa-public-2" {
  vpc_id            = aws_vpc.mwaa-vpc.id
  cidr_block        = var.public_subnet2_cidr
  availability_zone = var.region_az2
  map_public_ip_on_launch = true
  tags = {
    Name = "${var.mwaa_env_name} Public Subnet 2"
  }
}

# internet gateway for mwaa subnet
resource "aws_internet_gateway" "mwaa-inetgw" {
  vpc_id = aws_vpc.mwaa-vpc.id
  tags = {
    Name = var.mwaa_env_name
  }
}

# nat gateway for mwaa private subnet
resource "aws_eip" "mwaa-nat1" {
  vpc = true
  tags = {
    Name = "${var.mwaa_env_name} NAT Gateway 1"
  }
}
resource "aws_eip" "mwaa-nat2" {
  vpc = true
  tags = {
    Name = "${var.mwaa_env_name} NAT Gateway 2"
  }
}
resource "aws_nat_gateway" "mwaa-nat-gw1" {
  allocation_id = aws_eip.mwaa-nat1.id
  subnet_id     = aws_subnet.mwaa-public-1.id
  depends_on    = [aws_internet_gateway.mwaa-inetgw]
  tags = {
    Name = "${var.mwaa_env_name} NAT Gateway 1"
  }
}
resource "aws_nat_gateway" "mwaa-nat-gw2" {
  allocation_id = aws_eip.mwaa-nat2.id
  subnet_id     = aws_subnet.mwaa-public-2.id
  depends_on    = [aws_internet_gateway.mwaa-inetgw]
  tags = {
    Name = "${var.mwaa_env_name} NAT Gateway 2"
  }
}

# route table for mwaa public subnet
resource "aws_route_table" "mwaa-public-route" {
  vpc_id = aws_vpc.mwaa-vpc.id
  route {
    cidr_block = "0.0.0.0/0"
    gateway_id = aws_internet_gateway.mwaa-inetgw.id
  }
  tags = {
    Name = "${var.mwaa_env_name} Public Routes"
  }
}
resource "aws_route_table_association" "mwaa-public-route-inetgw1" {
  subnet_id      = aws_subnet.mwaa-public-1.id
  route_table_id = aws_route_table.mwaa-public-route.id
}
resource "aws_route_table_association" "mwaa-public-route-inetgw2" {
  subnet_id      = aws_subnet.mwaa-public-2.id
  route_table_id = aws_route_table.mwaa-public-route.id
}

# route table for mwaa private subnet
resource "aws_route_table" "mwaa-private-route1" {
  vpc_id = aws_vpc.mwaa-vpc.id
  route {
    cidr_block = "0.0.0.0/0"
    nat_gateway_id = aws_nat_gateway.mwaa-nat-gw1.id
  }
  tags = {
    Name = "${var.mwaa_env_name} Private Routes 1"
  }
}
resource "aws_route_table_association" "private-route1-assoc" {
  subnet_id      = aws_subnet.mwaa-private-1.id
  route_table_id = aws_route_table.mwaa-private-route1.id
}

resource "aws_route_table" "mwaa-private-route2" {
  vpc_id = aws_vpc.mwaa-vpc.id
  route {
    cidr_block = "0.0.0.0/0"
    nat_gateway_id = aws_nat_gateway.mwaa-nat-gw2.id
  }
  tags = {
    Name = "${var.mwaa_env_name} Private Routes 2"
  }
}
resource "aws_route_table_association" "private-route2-assoc" {
  subnet_id      = aws_subnet.mwaa-private-2.id
  route_table_id = aws_route_table.mwaa-private-route2.id
}

# iam role for mwaa
resource "aws_iam_role" "mwaa-execution" {
  name = "AmazonMWAA-${var.mwaa_env_name}"

  assume_role_policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Action": "sts:AssumeRole",
      "Principal": {
        "Service": [
          "airflow.amazonaws.com",
          "airflow-env.amazonaws.com"
        ]
      },
      "Effect": "Allow",
      "Sid": ""
    }
  ]
}
EOF
}

resource "aws_iam_role_policy" "mwaa-exec-policy" {
  name = "MWAA-Execution-Policy-${var.mwaa_env_name}"
  role = aws_iam_role.mwaa-execution.id

  policy = <<EOF
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": "airflow:PublishMetrics",
      "Resource": "arn:aws:airflow:${var.region}:${var.account_id}:environment/${var.mwaa_env_name}"
    },
    { 
      "Effect": "Deny",
      "Action": [ 
        "s3:ListAllMyBuckets"
      ],
      "Resource": [
        "${aws_s3_bucket.mwaa-scripts.arn}",
        "${aws_s3_bucket.mwaa-scripts.arn}/*"
      ]
    },
    { 
      "Effect": "Allow",
      "Action": [ 
        "s3:GetObject*",
        "s3:GetBucket*",
        "s3:List*"
      ],
      "Resource": [
        "${aws_s3_bucket.mwaa-scripts.arn}",
        "${aws_s3_bucket.mwaa-scripts.arn}/*"
      ]
    },
    {
      "Effect": "Allow",
      "Action": [
        "logs:CreateLogStream",
        "logs:CreateLogGroup",
        "logs:PutLogEvents",
        "logs:GetLogEvents",
        "logs:GetLogRecord",
        "logs:GetLogGroupFields",
        "logs:GetQueryResults"
      ],
      "Resource": [
        "arn:aws:logs:${var.region}:${var.account_id}:log-group:airflow-${var.mwaa_env_name}-*"
      ]
    },
    {
      "Effect": "Allow",
      "Action": [
        "logs:DescribeLogGroups"
      ],
      "Resource": [
        "*"
      ]
    },
    {
      "Effect": "Allow",
      "Action": "cloudwatch:PutMetricData",
      "Resource": "*"
    },
    {
      "Effect": "Allow",
      "Action": [
        "sqs:ChangeMessageVisibility",
        "sqs:DeleteMessage",
        "sqs:GetQueueAttributes",
        "sqs:GetQueueUrl",
        "sqs:ReceiveMessage",
        "sqs:SendMessage"
      ],
      "Resource": "arn:aws:sqs:${var.region}:*:airflow-celery-*"
    },
    {
      "Effect": "Allow",
      "Action": [
        "kms:Decrypt",
        "kms:DescribeKey",
        "kms:GenerateDataKey*",
        "kms:Encrypt"
      ],
      "NotResource": "arn:aws:kms:*:${var.account_id}:key/*",
      "Condition": {
        "StringLike": {
          "kms:ViaService": [
            "sqs.${var.region}.amazonaws.com"
          ]
        }
      }
    }
  ]
}
EOF
}

# security group for mwaa
resource "aws_security_group" "mwaa-execution" {
  name        = "airflow-security-group"
  description = "Security Group for ${var.mwaa_env_name}"
  vpc_id      = aws_vpc.mwaa-vpc.id

  ingress {
    from_port = 0
    to_port   = 0
    protocol  = -1
    self      = true
  }

  egress {
    from_port   = 0
    to_port     = 0
    protocol    = -1
    cidr_blocks = ["0.0.0.0/0"]
  }
}

# MWAA Environment
resource "aws_mwaa_environment" "mwaa-env" {
  source_bucket_arn = aws_s3_bucket.mwaa-scripts.arn
  dag_s3_path        = "dags"
  execution_role_arn = aws_iam_role.mwaa-execution.arn
  name               = var.mwaa_env_name
  max_workers        = 2
  webserver_access_mode = "PUBLIC_ONLY"
  network_configuration {
    security_group_ids = [aws_security_group.mwaa-execution.id]
    subnet_ids         = [aws_subnet.mwaa-private-1.id, aws_subnet.mwaa-private-2.id]
  }
  logging_configuration {
    task_logs {
      enabled   = true
      log_level = "INFO"
    }
    webserver_logs {
      enabled   = true
      log_level = "INFO"
    }
    scheduler_logs {
      enabled   = true
      log_level = "INFO"
    }
    worker_logs {
      enabled   = true
      log_level = "INFO"
    }
    dag_processing_logs {
      enabled   = true
      log_level = "INFO"
    }
  }
}