# Marker-driven join: completion markers landing in the output bucket invoke
# the join Lambda, which launches the fuse task when both extraction halves
# are complete, or holds the document and alerts when either half failed.
# The one-shot deadline schedules created by the trigger Lambda target the
# same function.

data "archive_file" "join" {
  type        = "zip"
  source_file = "${path.module}/../lambda/join.py"
  output_path = "${path.module}/build/join.zip"
}

resource "aws_lambda_function" "join" {
  function_name = "quber-join"
  role          = aws_iam_role.join.arn
  runtime       = "python3.13"
  handler       = "join.handler"
  timeout       = 60

  filename         = data.archive_file.join.output_path
  source_code_hash = data.archive_file.join.output_base64sha256

  environment {
    variables = {
      ECS_CLUSTER          = aws_ecs_cluster.quber.name
      TASK_DEFINITION      = aws_ecs_task_definition.table.family
      CONTAINER_NAME       = local.container_name
      SUBNETS              = join(",", data.aws_subnets.default.ids)
      SECURITY_GROUPS      = aws_security_group.task.id
      OUTPUT_BUCKET        = aws_s3_bucket.output.bucket
      INPUT_BUCKET         = data.aws_s3_bucket.input.bucket
      INPUT_PREFIX         = var.input_prefix
      ALERT_TOPIC_ARN      = aws_sns_topic.hold_alerts.arn
      RUNPOD_ENDPOINT_ID   = var.runpod_endpoint_id
      RUNPOD_API_KEY_PARAM = aws_ssm_parameter.runpod_api_key.name
    }
  }

  depends_on = [aws_cloudwatch_log_group.join]
}

resource "aws_cloudwatch_log_group" "join" {
  name              = "/aws/lambda/quber-join"
  retention_in_days = 30
}

resource "aws_lambda_permission" "output_s3_invoke" {
  statement_id   = "AllowOutputS3Invoke"
  action         = "lambda:InvokeFunction"
  function_name  = aws_lambda_function.join.function_name
  principal      = "s3.amazonaws.com"
  source_arn     = aws_s3_bucket.output.arn
  source_account = data.aws_caller_identity.current.account_id
}

# Marker events only: every job's last write is `*.complete.json`, so this
# suffix is the join signal. The handler itself ignores fuse markers and
# anything under hold/.
resource "aws_s3_bucket_notification" "output" {
  bucket = aws_s3_bucket.output.id

  lambda_function {
    lambda_function_arn = aws_lambda_function.join.arn
    events              = ["s3:ObjectCreated:*"]
    filter_suffix       = ".complete.json"
  }

  depends_on = [aws_lambda_permission.output_s3_invoke]
}

# Held documents page a human; subscription confirmed out of band.
resource "aws_sns_topic" "hold_alerts" {
  name = "quber-hold-alerts"
}

resource "aws_sns_topic_subscription" "hold_alerts_email" {
  count     = var.alert_email == "" ? 0 : 1
  topic_arn = aws_sns_topic.hold_alerts.arn
  protocol  = "email"
  endpoint  = var.alert_email
}

# ---------------------------------------------------------------------------
# Join Lambda role
# ---------------------------------------------------------------------------
resource "aws_iam_role" "join" {
  name               = "quber-join"
  assume_role_policy = data.aws_iam_policy_document.lambda_assume.json
}

data "aws_iam_policy_document" "join" {
  statement {
    sid       = "Logs"
    actions   = ["logs:CreateLogStream", "logs:PutLogEvents"]
    resources = ["${aws_cloudwatch_log_group.join.arn}:*"]
  }
  statement {
    sid = "MarkersHoldLock"
    # Read markers, create the fuse-started lock and hold objects. The lock
    # depends on conditional writes (If-None-Match), which ride PutObject.
    actions   = ["s3:GetObject", "s3:PutObject"]
    resources = ["${aws_s3_bucket.output.arn}/*"]
  }
  statement {
    sid = "ListOutput"
    # Without ListBucket, S3 answers a GET of a MISSING key with 403 instead
    # of 404, and the join treats that as an error rather than "marker not
    # written yet" — it crashed on every first-half marker event until this
    # was added.
    actions   = ["s3:ListBucket"]
    resources = [aws_s3_bucket.output.arn]
  }
  statement {
    sid       = "Alert"
    actions   = ["sns:Publish"]
    resources = [aws_sns_topic.hold_alerts.arn]
  }
  statement {
    sid       = "RunFuseTask"
    actions   = ["ecs:RunTask"]
    resources = ["arn:aws:ecs:${var.region}:${data.aws_caller_identity.current.account_id}:task-definition/${aws_ecs_task_definition.table.family}:*"]
  }
  statement {
    sid       = "PassTaskRoles"
    actions   = ["iam:PassRole"]
    resources = [aws_iam_role.task.arn, aws_iam_role.task_execution.arn]
  }
  statement {
    sid       = "ReadRunpodKey"
    actions   = ["ssm:GetParameter"]
    resources = [aws_ssm_parameter.runpod_api_key.arn]
  }
  statement {
    sid       = "DecryptParameters"
    actions   = ["kms:Decrypt"]
    resources = ["*"]
    condition {
      test     = "StringEquals"
      variable = "kms:ViaService"
      values   = ["ssm.${var.region}.amazonaws.com"]
    }
  }
}

resource "aws_iam_role_policy" "join" {
  name   = "join"
  role   = aws_iam_role.join.id
  policy = data.aws_iam_policy_document.join.json
}

# ---------------------------------------------------------------------------
# Scheduler role: lets the one-shot deadline schedules invoke the join Lambda
# ---------------------------------------------------------------------------
data "aws_iam_policy_document" "scheduler_assume" {
  statement {
    actions = ["sts:AssumeRole"]
    principals {
      type        = "Service"
      identifiers = ["scheduler.amazonaws.com"]
    }
  }
}

resource "aws_iam_role" "scheduler" {
  name               = "quber-parse-deadline-scheduler"
  assume_role_policy = data.aws_iam_policy_document.scheduler_assume.json
}

data "aws_iam_policy_document" "scheduler" {
  statement {
    sid       = "InvokeJoin"
    actions   = ["lambda:InvokeFunction"]
    resources = [aws_lambda_function.join.arn]
  }
}

resource "aws_iam_role_policy" "scheduler" {
  name   = "invoke-join"
  role   = aws_iam_role.scheduler.id
  policy = data.aws_iam_policy_document.scheduler.json
}