resource "aws_lambda_event_source_mapping" "sqs_fair_queue" {
  event_source_arn                   = var.source_sqs_arn
  function_name                      = aws_lambda_function.sqs_fair_queue.arn
  batch_size                         = 100
  maximum_batching_window_in_seconds = 30 # How long to wait to gather a batch, max: 300
}

# To install prereqs:
#   pip install --target ./scripts jsonpath-ng
data "archive_file" "sqs_fair_queue" {
  type        = "zip"
  source_dir  = "${path.module}/scripts/"
  output_path = "${path.module}/tmp/sqs_fair_queue.zip"
}


resource "aws_lambda_function" "sqs_fair_queue" {
  filename      = data.archive_file.sqs_fair_queue.output_path
  function_name = "sqs_fair_queue_${var.sqs_prefix}"
  role          = aws_iam_role.sqs_fair_queue.arn
  handler       = "sqs_fair_queue.lambda_handler"

  source_code_hash = data.archive_file.sqs_fair_queue.output_base64sha256

  runtime = "python3.9"

  environment {
    variables = {
      "SOURCE_SQS_ARN" = var.source_sqs_arn
      "SQS_PREFIX"     = var.sqs_prefix
      "NUM_QUEUES"     = var.num_queues
      "HASH_JSONPATH"  = var.hash_jsonpath
    }
  }
}

resource "aws_lambda_permission" "sqs_fair_queue" {
  statement_id  = "AllowExecutionFromSQS"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.sqs_fair_queue.function_name
  principal     = "sqs.amazonaws.com"
  source_arn    = var.source_sqs_arn
}

data "aws_iam_policy_document" "sqs_fair_queue" {
  statement {
    sid       = "SQSIngest"
    effect    = "Allow"
    resources = [var.source_sqs_arn]
    actions   = ["sqs:*"] # TODO: Nail down
    # Probably:
    #   "sqs:ReceiveMessage",
    #   "sqs:SendMessage",
    #   "sqs:GetQueueAttributes"
    #   "sqs:GetQueueUrl"
  }

  statement {
    sid       = "SQSPut"
    effect    = "Allow"
    resources = tolist(aws_sqs_queue.queue[*].arn)
    actions   = ["sqs:*"] # TODO: Nail down
  }
}

resource "aws_iam_policy" "sqs_fair_queue" {
  name        = "sqs_fair_queue_${var.sqs_prefix}"
  path        = "/sqs_fair_queue/"
  description = "SQS Fair Queueing Lambda Policy"
  policy      = data.aws_iam_policy_document.sqs_fair_queue.json
}

data "aws_iam_policy_document" "lambda_trust" {
  statement {
    sid     = ""
    effect  = "Allow"
    actions = ["sts:AssumeRole"]

    principals {
      type        = "Service"
      identifiers = ["lambda.amazonaws.com"]
    }
  }
}

resource "aws_iam_role" "sqs_fair_queue" {
  name               = "sqs_fair_queue_${var.sqs_prefix}"
  path               = "/sqs_fair_queue/"
  assume_role_policy = data.aws_iam_policy_document.lambda_trust.json
}

resource "aws_iam_role_policy_attachment" "sqs_fair_queue" {
  role       = aws_iam_role.sqs_fair_queue.name
  policy_arn = aws_iam_policy.sqs_fair_queue.arn
}

resource "aws_iam_role_policy_attachment" "aws_managed_lambda" {
  role       = aws_iam_role.sqs_fair_queue.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
}