provider "aws" {
  region = var.aws_region
}

terraform {
  backend "s3" {
    bucket  = "dev-orcd-terraform-state"
    key     = "dev/elisa-notebook/terraform.tfstate"
    region  = "us-east-1"
    encrypt = "true"
  }
}

data "aws_caller_identity" "current" {}

data "aws_vpc" "vpc" {
  tags = {
    Name = "dev"
  }
}

data "aws_subnets" "private" {
  filter {
    name   = "vpc-id"
    values = [data.aws_vpc.vpc.id]
  }

  tags = {
    Name = "private0*"
  }
}


data "aws_subnet" "private_subnets" {
  count = length(data.aws_subnets.private.ids)
  id    = element(local.subnet_ids_list, count.index)
}

locals {
  subnet_ids_string = join(",", data.aws_subnets.private.ids)
  subnet_ids_list   = split(",", local.subnet_ids_string)
  private_subnet_cidr_blocks = concat(
    local.https_listener_allow_cidr_blocks[var.environment],
    data.aws_subnet.private_subnets.*.cidr_block
  )
}


resource "aws_sagemaker_notebook_instance_lifecycle_configuration" "notebook_config" {
  name     = "${var.environment}-${var.service_name}-lifecycle-config"
  on_start = filebase64("on-start.sh")
}

resource "aws_kms_key" "sagemaker_kms_key" {
  description             = "${var.environment}-${var.service_name}"
  enable_key_rotation     = true
  deletion_window_in_days = 30
}

resource "aws_kms_alias" "sagemaker_kms_alias" {
  name          = "alias/${var.environment}-${var.service_name}"
  target_key_id = aws_kms_key.sagemaker_kms_key.key_id
}

resource "aws_security_group" "notebook_sg" {
  name        = "${var.environment}-${var.service_name}-security-group"
  description = "Notebooks security group for ${var.environment}-${var.service_name}"
  vpc_id      = data.aws_vpc.vpc.id
}

# RULES
# Security Group Rules
# allow https for sagemaker notebook (inbound)
resource "aws_security_group_rule" "allow_https_inbound" {
  type      = "ingress"
  from_port = var.https_listener_port
  to_port   = var.https_listener_port
  protocol  = "TCP"

  cidr_blocks = local.private_subnet_cidr_blocks

  security_group_id = aws_security_group.notebook_sg.id
}

# allow notebook to all CIDRs egress and all protocols
resource "aws_security_group_rule" "allow_notebook_egress" {
  type              = "egress"
  from_port         = 0
  to_port           = 0
  protocol          = "-1"
  cidr_blocks       = ["0.0.0.0/0"]
  security_group_id = aws_security_group.notebook_sg.id
}

data "aws_iam_policy_document" "sm_assume_role_policy" {
  statement {
    actions = ["sts:AssumeRole"]
    principals {
      type        = "Service"
      identifiers = ["sagemaker.amazonaws.com"]
    }
  }
}


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

    resources = var.lst_secret_arns

  }
  # do we need this?
  statement {
    actions = [
      "secretsmanager:GetRandomPassword",
    ]

    resources = [
      "*",
    ]
  }
}

resource "aws_iam_role" "notebook_iam_role" {
  name               = "${var.environment}-${var.service_name}-role"
  assume_role_policy = data.aws_iam_policy_document.sm_assume_role_policy.json
}

# Attaching the AWS default policy, "AmazonSageMakerFullAccess"
resource "aws_iam_policy_attachment" "sm_full_access_attach" {
  name       = "${var.environment}-${var.service_name}-policy-attachment"
  roles      = [aws_iam_role.notebook_iam_role.name]
  policy_arn = "arn:aws:iam::aws:policy/AmazonSageMakerFullAccess"
}

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


resource "aws_iam_role_policy_attachment" "sagemaker_snowflake_secrets_manager_policy_attachment" {
  role       = aws_iam_role.notebook_iam_role.name
  policy_arn = aws_iam_policy.sagemaker_secret_policy.arn
}

resource "aws_sagemaker_notebook_instance" "sagemaker_notebook" {
  name                   = "${var.environment}-${var.service_name}-notebook"
  role_arn               = aws_iam_role.notebook_iam_role.arn
  kms_key_id             = aws_kms_key.sagemaker_kms_key.key_id
  direct_internet_access = "Disabled"
  root_access            = "Disabled"
  security_groups        = [aws_security_group.notebook_sg.id]
  subnet_id              = element(tolist(data.aws_subnets.private.ids), 0)
  instance_type          = "ml.t3.medium"
  volume_size            = 5 // GB
  lifecycle_config_name  = aws_sagemaker_notebook_instance_lifecycle_configuration.notebook_config.name

  tags = {
    application_family = var.application_family
    environment        = var.environment
    terraform          = "true"
  }

}