# ---------------------------------------------------------------------------
# ECS task roles
# ---------------------------------------------------------------------------
data "aws_iam_policy_document" "ecs_assume" {
  statement {
    actions = ["sts:AssumeRole"]
    principals {
      type        = "Service"
      identifiers = ["ecs-tasks.amazonaws.com"]
    }
  }
}

# Execution role: pull the image, read the SSM parameters, write logs.
resource "aws_iam_role" "task_execution" {
  name               = "quber-task-execution"
  assume_role_policy = data.aws_iam_policy_document.ecs_assume.json
}

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

data "aws_iam_policy_document" "task_execution_secrets" {
  statement {
    sid     = "ReadParameters"
    actions = ["ssm:GetParameters"]
    resources = [
      aws_ssm_parameter.anthropic_api_key.arn,
      aws_ssm_parameter.logfire_token.arn,
      aws_ssm_parameter.langsmith_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" "task_execution_secrets" {
  name   = "ssm-read"
  role   = aws_iam_role.task_execution.id
  policy = data.aws_iam_policy_document.task_execution_secrets.json
}

# Task role: read inbound PDFs, write artifacts. Least privilege by prefix.
resource "aws_iam_role" "task" {
  name               = "quber-task"
  assume_role_policy = data.aws_iam_policy_document.ecs_assume.json
}

data "aws_iam_policy_document" "task_s3" {
  statement {
    sid       = "ReadInbound"
    actions   = ["s3:GetObject"]
    resources = ["${data.aws_s3_bucket.input.arn}/${var.input_prefix}*"]
  }
  statement {
    sid       = "ListInbound"
    actions   = ["s3:ListBucket"]
    resources = [data.aws_s3_bucket.input.arn]
    condition {
      test     = "StringLike"
      variable = "s3:prefix"
      values   = ["${var.input_prefix}*"]
    }
  }
  statement {
    sid = "WriteArtifacts"
    # GetObject as well as PutObject: cloudpathlib checks whether the key
    # already exists (HeadObject) before uploading each artifact.
    actions   = ["s3:PutObject", "s3:GetObject"]
    resources = ["${aws_s3_bucket.output.arn}/*"]
  }
  statement {
    sid = "ListOutput"
    # cloudpathlib's existence check falls back to ListObjects on the bucket
    # when HeadObject is denied, so the task needs ListBucket on the output.
    actions   = ["s3:ListBucket"]
    resources = [aws_s3_bucket.output.arn]
  }
}

resource "aws_iam_role_policy" "task_s3" {
  name   = "s3-access"
  role   = aws_iam_role.task.id
  policy = data.aws_iam_policy_document.task_s3.json
}

# ---------------------------------------------------------------------------
# Lambda trigger role
# ---------------------------------------------------------------------------
data "aws_iam_policy_document" "lambda_assume" {
  statement {
    actions = ["sts:AssumeRole"]
    principals {
      type        = "Service"
      identifiers = ["lambda.amazonaws.com"]
    }
  }
}

resource "aws_iam_role" "lambda" {
  name               = "quber-trigger"
  assume_role_policy = data.aws_iam_policy_document.lambda_assume.json
}

data "aws_iam_policy_document" "lambda" {
  statement {
    sid       = "Logs"
    actions   = ["logs:CreateLogStream", "logs:PutLogEvents"]
    resources = ["${aws_cloudwatch_log_group.lambda.arn}:*"]
  }
  statement {
    sid       = "ListOutput"
    actions   = ["s3:ListBucket"]
    resources = [aws_s3_bucket.output.arn]
  }
  statement {
    sid = "DedupAndPresignArtifacts"
    # GetObject backs the per-engine dedup HeadObject; PutObject backs the
    # presigned artifact/marker URLs handed to the RunPod worker — a
    # presigned URL executes with this role's permissions at use time.
    actions   = ["s3:GetObject", "s3:PutObject"]
    resources = ["${aws_s3_bucket.output.arn}/*"]
  }
  statement {
    sid = "PresignInboundRead"
    # Backs the presigned GET for the source PDF.
    actions   = ["s3:GetObject"]
    resources = ["${data.aws_s3_bucket.input.arn}/${var.input_prefix}*"]
  }
  statement {
    sid       = "RunTask"
    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       = "DecryptRunpodKey"
    actions   = ["kms:Decrypt"]
    resources = ["*"]
    condition {
      test     = "StringEquals"
      variable = "kms:ViaService"
      values   = ["ssm.${var.region}.amazonaws.com"]
    }
  }
  statement {
    sid       = "ScheduleDeadlineChecks"
    actions   = ["scheduler:CreateSchedule"]
    resources = ["arn:aws:scheduler:${var.region}:${data.aws_caller_identity.current.account_id}:schedule/default/quber-parse-deadline-*"]
  }
  statement {
    sid       = "PassSchedulerRole"
    actions   = ["iam:PassRole"]
    resources = [aws_iam_role.scheduler.arn]
  }
}

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