terraform {
  required_version = ">= 1.0"
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "~> 6.0"
    }
  }
}

provider "aws" {
  region = "ap-northeast-1"
}

data "aws_availability_zones" "available" {
  state = "available"
}

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

#--------------------------------------------------
# IAM: Operator Role
# LambdaがEC2インスタンスを管理するために必要なロール
#--------------------------------------------------

data "aws_iam_policy_document" "lambda_assume_role" {
  statement {
    effect = "Allow"

    principals {
      type        = "Service"
      identifiers = ["lambda.amazonaws.com"]
    }

    actions = ["sts:AssumeRole"]
  }
}

resource "aws_iam_role" "lambda_operator" {
  name               = "lambda-managed-instances-operator"
  assume_role_policy = data.aws_iam_policy_document.lambda_assume_role.json
}

# このポリシーにより、LambdaがEC2インスタンスのライフサイクルを管理可能になる
resource "aws_iam_role_policy_attachment" "lambda_operator" {
  role       = aws_iam_role.lambda_operator.name
  policy_arn = "arn:aws:iam::aws:policy/AWSLambdaManagedEC2ResourceOperator"
}

#--------------------------------------------------
# IAM: Lambda Execution Role
# Lambda関数の実行に必要なロール
#--------------------------------------------------

resource "aws_iam_role" "lambda_execution" {
  name               = "lambda-managed-instances-execution"
  assume_role_policy = data.aws_iam_policy_document.lambda_assume_role.json
}

resource "aws_iam_role_policy_attachment" "lambda_basic_execution" {
  role       = aws_iam_role.lambda_execution.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
}

resource "aws_iam_role_policy_attachment" "lambda_vpc_access" {
  role       = aws_iam_role.lambda_execution.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaVPCAccessExecutionRole"
}

#--------------------------------------------------
# VPC Resources
# Capacity ProviderにはVPC設定が必須
#--------------------------------------------------

resource "aws_vpc" "main" {
  cidr_block           = "10.0.0.0/16"
  enable_dns_hostnames = true
  enable_dns_support   = true

  tags = {
    Name = "lambda-managed-instances-vpc"
  }
}

# プライベートサブネット（2つのAZに分散配置）
resource "aws_subnet" "private" {
  count             = 2
  vpc_id            = aws_vpc.main.id
  cidr_block        = cidrsubnet(aws_vpc.main.cidr_block, 8, count.index)
  availability_zone = data.aws_availability_zones.available.names[count.index]

  tags = {
    Name = "lambda-managed-instances-private-${count.index + 1}"
  }
}

# セキュリティグループ（アウトバウンドのみ許可）
resource "aws_security_group" "lambda" {
  name        = "lambda-managed-instances-sg"
  description = "Security group for Lambda Managed Instances"
  vpc_id      = aws_vpc.main.id

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

  tags = {
    Name = "lambda-managed-instances-sg"
  }
}

#--------------------------------------------------
# Lambda Capacity Provider
# EC2インスタンスの配置先やインスタンス要件を定義
#--------------------------------------------------

resource "aws_lambda_capacity_provider" "main" {
  name = "example-capacity-provider"

  # VPC設定（必須）
  vpc_config {
    subnet_ids         = aws_subnet.private[*].id
    security_group_ids = [aws_security_group.lambda.id]
  }

  # Operator Roleの指定（必須）
  permissions_config {
    capacity_provider_operator_role_arn = aws_iam_role.lambda_operator.arn
  }

  # インスタンス要件（オプション）
  # 指定しない場合はLambdaが自動で最適なインスタンスを選択
  instance_requirements {
    architectures = ["x86_64"]
    # 特定のインスタンスタイプに限定する場合は以下を指定
    allowed_instance_types = ["m5.large"]
  }

  tags = {
    Name = "example-capacity-provider"
  }
}

#--------------------------------------------------
# Lambda Function
# capacity_provider_configを指定することでManaged Instancesとして動作
#--------------------------------------------------

resource "aws_lambda_function" "main" {
  filename         = data.archive_file.lambda.output_path
  function_name    = "example-managed-instances-function"
  role             = aws_iam_role.lambda_execution.arn
  handler          = "index.handler"
  source_code_hash = data.archive_file.lambda.output_base64sha256

  runtime     = "python3.13" # Managed InstancesではPython 3.13以降が必要
  memory_size = 2048         # Managed Instancesでは最低2 GB（2048 MB）が必要
  timeout     = 30

  # $LATEST.PUBLISHED形式のバージョンを作成する
  publish    = true
  publish_to = "LATEST_PUBLISHED"

  # このブロックを指定することでManaged Instancesとして動作する
  capacity_provider_config {
    lambda_managed_instances_capacity_provider_config {
      capacity_provider_arn = aws_lambda_capacity_provider.main.arn
    }
  }

  tags = {
    Name = "example-managed-instances-function"
  }

  depends_on = [
    aws_iam_role_policy_attachment.lambda_basic_execution,
    aws_iam_role_policy_attachment.lambda_vpc_access,
  ]
}

import json

def handler(event, context):
    return {
        "statusCode": 200,
        "body": json.dumps({
            "message": "Hello from Lambda Managed Instances!"
        })
    }

% aws lambda invoke --function-name example-managed-instances-function outfile && cat outfile 

{
    "StatusCode": 200,
    "ExecutedVersion": "$LATEST.PUBLISHED"
}
{"statusCode": 200, "body": "{\"message\": \"Hello from Lambda Managed Instances!\"}"}

aws_lambda_capacity_provider.main: Still destroying... [name=example-capacity-provider, 10m0s elapsed]
aws_lambda_capacity_provider.main: Destruction complete after 10m5s