data "aws_s3_bucket" "selected" {
  bucket = var.project_bucket
}

data "aws_iam_policy_document" "ds_moments_pipeline_job_assume_role_policy_document" {
  statement {
    actions = ["sts:AssumeRole"]

    principals {
      type = "Service"
      identifiers = [
        "ecs.amazonaws.com",
        "ec2.amazonaws.com",
        "ecs-tasks.amazonaws.com"
      ]
    }
  }
}

resource "aws_iam_role" "ds_moments_pipeline_job_role" {
  name               = "${var.environment}-${var.service_name}-job-role"
  assume_role_policy = data.aws_iam_policy_document.ds_moments_pipeline_job_assume_role_policy_document.json
}

# Create policy to access S3 bucket from Batch jobs
data "aws_iam_policy_document" "ds_moments_pipeline_s3_policy_document" {
  statement {
    actions = [
      "s3:GetObject",
      "s3:PutObject",
    ]
    resources = [
      "${data.aws_s3_bucket.selected.arn}/*"
    ]
  }

  statement {
    actions = [
      "s3:GetBucketLocation",
      "s3:ListBucket",
    ]
    resources = [
      data.aws_s3_bucket.selected.arn
    ]
  }
}

resource "aws_iam_policy" "ds_moments_pipeline_s3_policy" {
  name        = "S3-${var.environment}-${var.service_name}-RW-policy"
  description = "RW access to the ${var.project_bucket} bucket"
  policy      = data.aws_iam_policy_document.ds_moments_pipeline_s3_policy_document.json
}

resource "aws_iam_role_policy_attachment" "streamlit_audience_job_s3_bucket_policy_attachment" {
  role       = aws_iam_role.ds_moments_pipeline_job_role.name
  policy_arn = aws_iam_policy.ds_moments_pipeline_s3_policy.arn
}

data "aws_iam_policy_document" "secrets_manager_policy" {
  statement {
    actions = [
      "secretsmanager:GetResourcePolicy",
      "secretsmanager:GetSecretValue",
      "secretsmanager:DescribeSecret",
      "secretsmanager:ListSecretVersionIds",
    ]

    resources = [
      "arn:aws:secretsmanager:*:*:secret:dev/sagemaker-notebook-instance/",
      "arn:aws:secretsmanager:*:*:secret:dev/sagemaker-notebook-instance/*",
    ]
  }

  statement {
    actions = [
      "secretsmanager:GetRandomPassword",
    ]

    resources = [
      "*",
    ]
  }
}

# Secrets manager policy for job role
resource "aws_iam_policy" "ds_moments_pipeline_secret_policy" {
  name   = "SecretsManager-${var.environment}-${var.service_name}-policy"
  policy = data.aws_iam_policy_document.secrets_manager_policy.json
}

# Attach secrets manager policy to task role
resource "aws_iam_role_policy_attachment" "ds_moments_job_secrets_manager_policy_attachment" {
  role       = aws_iam_role.ds_moments_pipeline_job_role.name
  policy_arn = aws_iam_policy.ds_moments_pipeline_secret_policy.arn
}

resource "aws_iam_role_policy_attachment" "ds_moments_job_ecs_policy_attachment" {
  role       = aws_iam_role.ds_moments_pipeline_job_role.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy"
}

# Definition for job nr 1
resource "aws_batch_job_definition" "ds_moments_pipeline_1_job_definition" {
  name = "${var.environment}-${var.service_name}-1-job-definition"
  type = "container"
  platform_capabilities = [
    "FARGATE",
  ]

  timeout {
    attempt_duration_seconds = 900 # 15 minutes
  }

  tags = local.tags

  container_properties = <<CONTAINER_PROPERTIES
{
  "image": "${data.aws_caller_identity.current.account_id}.dkr.ecr.${var.aws_region}.amazonaws.com/${var.github_repository}:latest",
  "fargatePlatformConfiguration": {
    "platformVersion": "LATEST"
  },
  "resourceRequirements": [
    {"type": "VCPU", "value": "1"},
    {"type": "MEMORY", "value": "3072"}
  ],
  "executionRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "jobRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "environment": [
    {"name": "STEP", "value": "${var.step_1_name}"},
    {"name": "CHUNK_COUNT", "value": "${var.chunk_count}"},
    {"name": "FOLDER", "value": "${var.folder}"},
    {"name": "RUN_ID", "value": "${var.run_id}"}
  ]
}
CONTAINER_PROPERTIES
}

# Definition for job nr 2
resource "aws_batch_job_definition" "ds_moments_pipeline_2_job_definition" {
  name = "${var.environment}-${var.service_name}-2-job-definition"
  type = "container"
  platform_capabilities = [
    "FARGATE",
  ]

  timeout {
    attempt_duration_seconds = 1800 # 30 minutes
  }

  tags = local.tags

  container_properties = <<CONTAINER_PROPERTIES
{
  "image": "${data.aws_caller_identity.current.account_id}.dkr.ecr.${var.aws_region}.amazonaws.com/${var.github_repository}:latest",
  "fargatePlatformConfiguration": {
    "platformVersion": "LATEST"
  },
  "resourceRequirements": [
    {"type": "VCPU", "value": "2"},
    {"type": "MEMORY", "value": "4096"}
  ],
  "executionRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "jobRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "environment": [
    {"name": "STEP", "value": "${var.step_2_name}"},
    {"name": "CHUNK_COUNT", "value": "${var.chunk_count}"},
    {"name": "FOLDER", "value": "${var.folder}"},
    {"name": "RUN_ID", "value": "${var.run_id}"}
  ]
}
CONTAINER_PROPERTIES
}

# Definition for job nr 3
resource "aws_batch_job_definition" "ds_moments_pipeline_3_job_definition" {
  name = "${var.environment}-${var.service_name}-3-job-definition"
  type = "container"
  platform_capabilities = [
    "FARGATE",
  ]

  timeout {
    attempt_duration_seconds = 5400 # 90 minutes
  }

  tags = local.tags

  container_properties = <<CONTAINER_PROPERTIES
{
  "image": "${data.aws_caller_identity.current.account_id}.dkr.ecr.${var.aws_region}.amazonaws.com/${var.github_repository}:latest",
  "fargatePlatformConfiguration": {
    "platformVersion": "LATEST"
  },
  "resourceRequirements": [
    {"type": "VCPU", "value": "8"},
    {"type": "MEMORY", "value": "16384"}
  ],
  "executionRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "jobRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "environment": [
    {"name": "STEP", "value": "${var.step_3_name}"},
    {"name": "FOLDER", "value": "${var.folder}"},
    {"name": "RUN_ID", "value": "${var.run_id}"}
  ]
}
CONTAINER_PROPERTIES
}

# Definition for job nr 4
resource "aws_batch_job_definition" "ds_moments_pipeline_4_job_definition" {
  name = "${var.environment}-${var.service_name}-4-job-definition"
  type = "container"
  platform_capabilities = [
    "FARGATE",
  ]

  timeout {
    attempt_duration_seconds = 900 # 15 minutes
  }

  tags = local.tags

  container_properties = <<CONTAINER_PROPERTIES
{
  "image": "${data.aws_caller_identity.current.account_id}.dkr.ecr.${var.aws_region}.amazonaws.com/${var.github_repository}:latest",
  "fargatePlatformConfiguration": {
    "platformVersion": "LATEST"
  },
  "resourceRequirements": [
    {"type": "VCPU", "value": "1"},
    {"type": "MEMORY", "value": "3072"}
  ],
  "executionRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "jobRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "environment": [
    {"name": "STEP", "value": "${var.step_4_name}"},
    {"name": "FOLDER", "value": "${var.folder}"},
    {"name": "RUN_ID", "value": "${var.run_id}"}
  ]
}
CONTAINER_PROPERTIES
}

# Definition for job nr 5
resource "aws_batch_job_definition" "ds_moments_pipeline_5_job_definition" {
  name = "${var.environment}-${var.service_name}-5-job-definition"
  type = "container"
  platform_capabilities = [
    "FARGATE",
  ]

  timeout {
    attempt_duration_seconds = 900 # 15 minutes
  }

  tags = local.tags

  container_properties = <<CONTAINER_PROPERTIES
{
  "image": "${data.aws_caller_identity.current.account_id}.dkr.ecr.${var.aws_region}.amazonaws.com/${var.github_repository}:latest",
  "fargatePlatformConfiguration": {
    "platformVersion": "LATEST"
  },
  "resourceRequirements": [
    {"type": "VCPU", "value": "1"},
    {"type": "MEMORY", "value": "3072"}
  ],
  "executionRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "jobRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "environment": [
    {"name": "STEP", "value": "${var.step_5_name}"},
    {"name": "FOLDER", "value": "${var.folder}"},
    {"name": "RUN_ID", "value": "${var.run_id}"}
  ]
}
CONTAINER_PROPERTIES
}

# Definition for job nr 6
resource "aws_batch_job_definition" "ds_moments_pipeline_6_job_definition" {
  name = "${var.environment}-${var.service_name}-6-job-definition"
  type = "container"
  platform_capabilities = [
    "FARGATE",
  ]

  timeout {
    attempt_duration_seconds = 900 # 15 minutes
  }

  tags = local.tags

  container_properties = <<CONTAINER_PROPERTIES
{
  "image": "${data.aws_caller_identity.current.account_id}.dkr.ecr.${var.aws_region}.amazonaws.com/${var.github_repository}:latest",
  "fargatePlatformConfiguration": {
    "platformVersion": "LATEST"
  },
  "resourceRequirements": [
    {"type": "VCPU", "value": "2"},
    {"type": "MEMORY", "value": "4096"}
  ],
  "executionRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "jobRoleArn": "${aws_iam_role.ds_moments_pipeline_job_role.arn}",
  "environment": [
    {"name": "STEP", "value": "${var.step_6_name}"},
    {"name": "FOLDER", "value": "${var.folder}"},
    {"name": "RUN_ID", "value": "${var.run_id}"}
  ]
}
CONTAINER_PROPERTIES
}