provider "aws" {
  region = var.region 
}

################################################################################
# create VPC 
################################################################################

resource "random_id" "random" {
  byte_length = 8
}

data "aws_availability_zones" "available" {
}

module "vpc" {
  source = "terraform-aws-modules/vpc/aws"
  version = "3.7.0"

  name = "tidb_demo_vpc${random_id.random.hex}"
  cidr = "10.0.0.0/16"

  azs  = data.aws_availability_zones.available.names 
  private_subnets = ["10.0.1.0/24", "10.0.2.0/24", "10.0.3.0/24"]
  public_subnets  = ["10.0.101.0/24", "10.0.102.0/24", "10.0.103.0/24"]


  enable_nat_gateway = true
  single_nat_gateway = true
  enable_dns_hostnames = true

  public_subnet_tags = {
    Type = "public"
  }

  tags = {
    Terraform = "true"
    Environment = "demo"
  }
}

resource "aws_vpc_endpoint" "s3" {
  vpc_id = module.vpc.vpc_id
  service_name = "com.amazonaws.${var.region}.s3"
  route_table_ids = concat(module.vpc.public_route_table_ids, module.vpc.private_route_table_ids)
}


################################################################################
# shared resources
################################################################################

resource "aws_security_group" "tidb_cluster" {
  name        = "tidb_cluster"
  description = "Allow all tidb cluster component to talk to each other" 
  vpc_id      = module.vpc.vpc_id 
}

resource "aws_security_group_rule" "tidb_cluster_in" {
  type              = "ingress"
  from_port         = 0 
  to_port           = 0
  protocol          = "-1"
  self              = true
  security_group_id = aws_security_group.tidb_cluster.id 
}

resource "aws_security_group_rule" "tidb_cluster_out" {
  type              = "egress"
  from_port         = 0 
  to_port           = 0
  protocol          = "-1"
  cidr_blocks       = ["0.0.0.0/0"] 
  ipv6_cidr_blocks  = ["::/0"]
  security_group_id = aws_security_group.tidb_cluster.id 
}

# key pair
################################################################################
resource "tls_private_key" "demo" {
  algorithm = "RSA"
}

resource "aws_key_pair" "demo" {
  key_name   = "demo${random_id.random.hex}"
  public_key = tls_private_key.demo.public_key_openssh

  lifecycle {
    ignore_changes = [key_name]
  }
}

resource "local_file" "public_key" {
  content  = tls_private_key.demo.public_key_openssh
  filename = "${path.module}/tidb_demo${random_id.random.hex}.pub"
}

resource "local_file" "private_key" {
  content  = tls_private_key.demo.private_key_pem
  filename = "${path.module}/tidb_demo${random_id.random.hex}.pem"
}

resource "null_resource" "chmod" {
  depends_on = [local_file.private_key]

  triggers = {
    key = tls_private_key.demo.private_key_pem
  }

  provisioner "local-exec" {
    command = "chmod 400 ${local_file.private_key.filename}"
  }
}

################################################################################
# jump server
################################################################################

resource "aws_security_group" "jump" {
  name        = "jump"
  description = "Allow ssh access to jump server" 
  vpc_id      = module.vpc.vpc_id 
}

resource "aws_security_group_rule" "ssh" {
  type              = "ingress"
  from_port         = 22 
  to_port           = 22
  protocol          = "tcp"
  cidr_blocks       = ["0.0.0.0/0"] 
  security_group_id = aws_security_group.jump.id 
}

resource "aws_security_group_rule" "jump_out" {
  type              = "egress"
  from_port         = 0 
  to_port           = 0
  protocol          = "-1"
  cidr_blocks       = ["0.0.0.0/0"] 
  ipv6_cidr_blocks  = ["::/0"]
  security_group_id = aws_security_group.jump.id 
}

# instance profile to access s3 bucket
resource "aws_iam_instance_profile" "jump" {
  name = "jump${random_id.random.hex}"
  role = aws_iam_role.jump.name
}

resource "aws_iam_role" "jump" {
  name = "jump${random_id.random.hex}"

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

  inline_policy {
    name = "my_inline_policy"

    policy = jsonencode({
      Version = "2012-10-17"
      Statement = [
        {
          Action   = ["s3:*"]
          Effect   = "Allow"
          Resource = [aws_s3_bucket.tidb_demo.arn, "${aws_s3_bucket.tidb_demo.arn}/*"]
        },
      ]
    })
  }
}

resource "aws_instance" "jump" {
  ami = var.ami_id[var.region] 
  instance_type = var.jump_instance_type 
  subnet_id = module.vpc.public_subnets[0]

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

  vpc_security_group_ids = [aws_security_group.tidb_cluster.id, aws_security_group.jump.id]
  key_name = aws_key_pair.demo.key_name
  iam_instance_profile = aws_iam_instance_profile.jump.name

  user_data = templatefile("${path.module}/jump_userdata.yaml",{
                bucket_name = aws_s3_bucket.tidb_demo.bucket
                private_key = filebase64(local_file.private_key.filename)
              })

  tags = {
    Name = "jump server"
  }

  # wait until all init commands finishes
  provisioner "remote-exec" {
    inline = [
      "sudo cloud-init status --wait"
    ]
   
    connection {
        type     = "ssh"
        user     = "centos"
        host     = aws_instance.jump.public_ip
        private_key = file(local_file.private_key.filename)
    }
  }

  # jump server needs to wait for all s3 files been uploaded
  depends_on = [
    aws_s3_bucket_object.files,
    aws_instance.pd,
    aws_instance.tidb,
    aws_instance.tikv,
  ]
}


################################################################################
# create tidb servers 
################################################################################

# pd
################################################################################
resource "aws_instance" "pd" {
  count = var.pd_instance_number

  ami = var.ami_id[var.region] 
  instance_type = var.pd_instance_type 
  
  subnet_id = module.vpc.public_subnets[count.index % length(module.vpc.public_subnets)]

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

  vpc_security_group_ids = [aws_security_group.tidb_cluster.id, aws_security_group.tidb_monitoring.id]
  key_name = aws_key_pair.demo.key_name

  tags = {
    Name = "PD ${count.index}"
  }
}

# tidb
################################################################################
resource "aws_instance" "tidb" {
  count = var.tidb_instance_number + var.tidb_instance_number_scaleout 

  ami = var.ami_id[var.region] 
  instance_type = var.tidb_instance_type 
  
  subnet_id = module.vpc.private_subnets[count.index % length(module.vpc.private_subnets)]

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

  vpc_security_group_ids = [aws_security_group.tidb_cluster.id, aws_security_group.tidb_nlb.id]
  key_name = aws_key_pair.demo.key_name

  tags = {
    Name = "TiDB ${count.index}"
  }
}

# nlb for tidb 
################################################################################
resource "aws_lb" "tidb_lb" {
  name               = "tidb-lb${random_id.random.hex}"
  internal           = true
  load_balancer_type = "network"
  subnets            = module.vpc.private_subnets 
  enable_cross_zone_load_balancing = true
}

resource "aws_lb_target_group" "tidb_lb" {
  name     = "tidb-lb${random_id.random.hex}"
  port     = 4000
  protocol = "TCP"
  vpc_id   = module.vpc.vpc_id 
  target_type = "instance"

  health_check {
    port = 4000
    protocol = "TCP"
    healthy_threshold = 2 
    unhealthy_threshold = 2 
  }
}

resource "aws_lb_listener" "tidb_lb" {
  load_balancer_arn = aws_lb.tidb_lb.arn
  port              = "4000"
  protocol          = "TCP"

  default_action {
    type             = "forward"
    target_group_arn = aws_lb_target_group.tidb_lb.arn
  }
}

resource "aws_lb_target_group_attachment" "tidb" {
  count = length(aws_instance.tidb)
  target_group_arn = aws_lb_target_group.tidb_lb.arn
  target_id        = aws_instance.tidb[count.index].id 
  port             = 4000
}

data "aws_network_interface" "tidb_lb" {
  count = length(module.vpc.private_subnets)

  filter {
    name   = "description"
    values = ["ELB ${aws_lb.tidb_lb.arn_suffix}"]
  }

  filter {
    name   = "subnet-id"
    values = [module.vpc.private_subnets[count.index]]
  }
}

resource "aws_security_group" "tidb_nlb" {
  name        = "tidb-nlb"
  description = "allow mysql connection from nlb as well as jump server" 
  vpc_id      = module.vpc.vpc_id 
}

resource "aws_security_group_rule" "tidb_nlb_in" {
  type              = "ingress"
  from_port         = 4000 
  to_port           = 4000
  protocol          = "TCP"
  cidr_blocks       = formatlist("%s/32", concat([for eni in data.aws_network_interface.tidb_lb : eni.private_ip], [aws_instance.jump.private_ip]))
  security_group_id = aws_security_group.tidb_nlb.id 
}

# tikv 
################################################################################
resource "aws_instance" "tikv" {
  count = var.tikv_instance_number + var.tikv_instance_number_scaleout 

  ami = var.ami_id[var.region] 
  instance_type = var.tikv_instance_type 
  
  subnet_id = module.vpc.private_subnets[count.index % length(module.vpc.private_subnets)]
  iam_instance_profile = aws_iam_instance_profile.jump.name

  root_block_device {
    volume_size = var.tikv_instance_disk_size
    volume_type = "gp3"
    iops = var.tikv_ebs_iops
    throughput = var.tikv_ebs_throughput
  }

  vpc_security_group_ids = [aws_security_group.tidb_cluster.id]
  key_name = aws_key_pair.demo.key_name

  tags = {
    Name = "TiKV ${count.index}"
  }
}

# monitoring 
################################################################################
resource "aws_security_group" "tidb_monitoring" {
  name        = "tidb_monitoring"
  description = "Allow visit to tidb monitoring publicly" 
  vpc_id      = module.vpc.vpc_id 
}

resource "aws_security_group_rule" "tidb_monitoring_dashboard_in" {
  type              = "ingress"
  from_port         = 2379 
  to_port           = 2379
  protocol          = "tcp"
  cidr_blocks       = ["0.0.0.0/0"] 
  security_group_id = aws_security_group.tidb_monitoring.id 
}

resource "aws_security_group_rule" "tidb_monitoring_grafana_in" {
  type              = "ingress"
  from_port         = 3000 
  to_port           = 3000
  protocol          = "tcp"
  cidr_blocks       = ["0.0.0.0/0"] 
  security_group_id = aws_security_group.tidb_monitoring.id 
}

resource "aws_security_group_rule" "tidb_monitoring_out" {
  type              = "egress"
  from_port         = 0 
  to_port           = 0
  protocol          = "-1"
  cidr_blocks       = ["0.0.0.0/0"] 
  ipv6_cidr_blocks  = ["::/0"]
  security_group_id = aws_security_group.tidb_monitoring.id 
}

# s3 bucket and files to store supporting data of tidb cluster creation 
################################################################################

resource "aws_s3_bucket" "tidb_demo" {
  bucket = "tidb-demo${random_id.random.hex}"
  acl    = "private"
}

resource "aws_s3_bucket_object" "files" {
  for_each = fileset("${path.module}/..", "{demos,tiup}/**/files/*.{template,sh}")

  key     = each.value 
  bucket  = aws_s3_bucket.tidb_demo.id
  acl     = "private"

  content = templatefile("${path.module}/../${each.value}", {
                pd_ips = aws_instance.pd.*.private_ip
                monitor_ips = aws_instance.pd.*.private_ip
                dashboard_ip = aws_instance.pd[0].private_ip
                tidb_ips = slice(aws_instance.tidb.*.private_ip, 0, var.tidb_instance_number)
                tidb_scale_ips = slice(aws_instance.tidb.*.private_ip, var.tidb_instance_number, var.tidb_instance_number + var.tidb_instance_number_scaleout) 
                tikv_ips = slice(aws_instance.tikv.*.private_ip, 0, var.tikv_instance_number)
                tikv_scale_ips = slice(aws_instance.tikv.*.private_ip, var.tikv_instance_number, var.tikv_instance_number + var.tikv_instance_number_scaleout) 
                s3_bucket = aws_s3_bucket.tidb_demo.id
                region = var.region
                tidb_lb = aws_lb.tidb_lb.dns_name
                tiup_config_path = var.tiup_config_path
            }) 
}