data "aws_vpc" "default" {
  count   = var.vpc_id == null ? 1 : 0
  default = true
}

data "aws_subnets" "default" {
  count = var.subnet_id == null ? 1 : 0

  filter {
    name   = "vpc-id"
    values = [local.vpc_id]
  }
}

data "aws_subnet" "selected" {
  id = local.subnet_id
}

data "aws_ami" "nixos" {
  count       = var.base_nixos_ami_id == null ? 1 : 0
  most_recent = true
  owners      = var.base_nixos_ami_owners

  filter {
    name   = "name"
    values = [local.base_nixos_ami_name_pattern]
  }

  filter {
    name   = "architecture"
    values = [local.aws_architecture]
  }

  filter {
    name   = "virtualization-type"
    values = ["hvm"]
  }
}

locals {
  vpc_id           = var.vpc_id == null ? data.aws_vpc.default[0].id : var.vpc_id
  subnet_id        = var.subnet_id == null ? data.aws_subnets.default[0].ids[0] : var.subnet_id
  aws_architecture = var.system == "aarch64-linux" ? "arm64" : "x86_64"
  base_nixos_ami_name_pattern = (
    var.base_nixos_ami_name_pattern == null
    ? "nixos/*-${var.system}"
    : var.base_nixos_ami_name_pattern
  )
  base_nixos_ami_id = var.base_nixos_ami_id == null ? data.aws_ami.nixos[0].id : var.base_nixos_ami_id
  resource_name     = "${var.name}-builder-${local.aws_architecture}-${substr(var.input_hash, 0, 12)}"
  nix_device_id     = replace(aws_ebs_volume.nix.id, "-", "")
  nix_device        = "/dev/disk/by-id/nvme-Amazon_Elastic_Block_Store_${local.nix_device_id}"
  tags = {
    Name          = local.resource_name
    Project       = var.name
    Role          = "ami-builder"
    Architecture  = local.aws_architecture
    System        = var.system
    InputHash     = var.input_hash
    BuilderGitRev = var.builder_git_rev
    FlakeLockRev  = var.flake_lock_rev
  }
}

resource "aws_iam_role" "builder" {
  name = "${local.resource_name}-instance"

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

resource "aws_iam_role_policy_attachment" "ssm" {
  role       = aws_iam_role.builder.name
  policy_arn = "arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore"
}

resource "aws_iam_instance_profile" "builder" {
  name = "${local.resource_name}-instance"
  role = aws_iam_role.builder.name
}

resource "aws_security_group" "builder" {
  name        = "${local.resource_name}-instance"
  description = "Temporary NixOS ZFS AMI builder egress"
  vpc_id      = local.vpc_id

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

  tags = local.tags
}

resource "aws_key_pair" "admin" {
  count      = var.ssh_public_key == null ? 0 : 1
  key_name   = "${local.resource_name}-admin"
  public_key = var.ssh_public_key
}

resource "aws_ebs_volume" "nix" {
  availability_zone = data.aws_subnet.selected.availability_zone
  size              = var.nix_volume_size_gb
  type              = "gp3"

  tags = merge(local.tags, {
    Name = "${local.resource_name}-nix"
  })
}

resource "aws_instance" "builder" {
  ami                         = local.base_nixos_ami_id
  instance_type               = var.builder_instance_type
  subnet_id                   = local.subnet_id
  vpc_security_group_ids      = [aws_security_group.builder.id]
  iam_instance_profile        = aws_iam_instance_profile.builder.name
  associate_public_ip_address = true
  key_name                    = var.ssh_public_key == null ? null : aws_key_pair.admin[0].key_name
  user_data_replace_on_change = true

  user_data = templatefile("${path.module}/user-data.sh.tftpl", {
    builder_git_rev = var.builder_git_rev
    flake_lock_rev  = var.flake_lock_rev
    input_hash      = var.input_hash
    nix_device      = local.nix_device
    nix_pool        = var.nix_pool
    ssh_keys        = jsonencode(compact([var.ssh_public_key]))
    system          = var.system
  })

  root_block_device {
    volume_size = var.root_volume_size_gb
    volume_type = "gp3"
  }

  tags = local.tags
}

resource "aws_volume_attachment" "nix" {
  device_name = var.nix_device_name
  volume_id   = aws_ebs_volume.nix.id
  instance_id = aws_instance.builder.id
}