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

provider "aws" {
  region = var.aws_region
}

# Caller Identity
data "aws_caller_identity" "current" {}

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

# Private Subnets
data "aws_subnet_ids" "private" {
  vpc_id = data.aws_vpc.vpc.id
  tags = {
    Name = "private0*"
  }
}

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

# locals
locals {
  subnet_ids_string = join(",", data.aws_subnet_ids.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
  )
}


# notebook kms key
data "aws_kms_key" "sagemaker_kms_key" {
  key_id = var.kms_arn
}

# SECURITY GROUP
# define security group for notebook instance
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" {
  count     = var.task_type == "web_service" ? var.https_listener_enabled ? 1 : 0 : 0
  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
}

# SAGE MAKER NOTEBOOK
# for training/crunching bigger dataframes (training models that require gpu) 
# Note: Not a 24/7 instance, runs 2-6 hours a week
resource "aws_sagemaker_notebook_instance" "orchard_gpu_notebook_light" {
  name                   = "${var.environment}-orchard-gpu-notebook-light"
  role_arn               = var.sagemaker_execution_role_arn
  kms_key_id             = data.aws_kms_key.sagemaker_kms_key.id
  direct_internet_access = "Disabled"
  root_access            = "Disabled"
  security_groups        = [aws_security_group.notebook_sg.id]
  subnet_id              = element(tolist(data.aws_subnet_ids.private.ids), 0)
  instance_type          = "ml.g4dn.2xlarge" //most cost-effective GPU option to run
  volume_size            = 50                // GB

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

}