############################################
# Provider
############################################
provider "aws" {
  region  = var.region
  profile = "DBaaS-DevUser-Role-219248915861"
}

############################################
# Create a Fresh VPC
############################################
resource "aws_vpc" "tidb_vpc" {
  cidr_block = "10.50.0.0/16"
  tags = {
    Name = "TiDB-VPC"
  }
}

############################################
# Internet Gateway & Public Route
############################################
resource "aws_internet_gateway" "igw" {
  vpc_id = aws_vpc.tidb_vpc.id
  tags = {
    Name = "TiDB-IGW"
  }
}

resource "aws_route_table" "public_rt" {
  vpc_id = aws_vpc.tidb_vpc.id
  tags = {
    Name = "TiDB-Public-RouteTable"
  }
}

resource "aws_route" "default" {
  route_table_id         = aws_route_table.public_rt.id
  destination_cidr_block = "0.0.0.0/0"
  gateway_id             = aws_internet_gateway.igw.id
}

############################################
# Subnets
############################################
resource "aws_subnet" "public_subnet" {
  count = var.num_clusters

  vpc_id                  = aws_vpc.tidb_vpc.id
  cidr_block              = "10.50.${count.index}.0/25"
  availability_zone       = "${var.region}a"
  map_public_ip_on_launch = true

  tags = {
    Name = "TiDB-Public-Subnet-${count.index}"
  }
}

resource "aws_subnet" "private_subnet" {
  count = var.num_clusters

  vpc_id                  = aws_vpc.tidb_vpc.id
  cidr_block              = "10.50.${count.index}.128/25"
  availability_zone       = "${var.region}a"
  map_public_ip_on_launch = false

  tags = {
    Name = "TiDB-Private-Subnet-${count.index}"
  }
}

resource "aws_route_table_association" "public_assoc" {
  count = var.num_clusters

  subnet_id      = aws_subnet.public_subnet[count.index].id
  route_table_id = aws_route_table.public_rt.id
}

############################################
# Security Group
############################################
resource "aws_security_group" "tidb_sg" {
  name   = "TiDB-SG"
  vpc_id = aws_vpc.tidb_vpc.id

  ingress {
    description = "SSH"
    from_port   = 22
    to_port     = 22
    protocol    = "tcp"
    cidr_blocks = ["0.0.0.0/0"]
  }

  ingress {
    description = "TiDB SQL"
    from_port   = 4000
    to_port     = 4000
    protocol    = "tcp"
    cidr_blocks = ["0.0.0.0/0"]
  }

  ingress {
    description = "All cluster internal"
    from_port   = 0
    to_port     = 65535
    protocol    = "tcp"
    cidr_blocks = ["0.0.0.0/0"]
  }

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

  tags = {
    Name = "TiDB-SG"
  }
}

############################################
# Bastion Hosts
############################################
resource "aws_instance" "bastion" {
  count         = var.num_clusters
  ami           = var.ami_id
  instance_type = "m5.4xlarge"

  subnet_id              = aws_subnet.public_subnet[count.index].id
  vpc_security_group_ids = [aws_security_group.tidb_sg.id]
  key_name               = "tidb-demo"

  tags = {
    Name = "Bastion-Host-${count.index}"
  }
}

############################################
# TiDB, TiKV, PD, TiFlash Instances
############################################
resource "aws_instance" "tidb" {
  count         = 2 * var.num_clusters
  ami           = var.ami_id
  instance_type = var.tidb_instance_type

  subnet_id              = element(aws_subnet.private_subnet[*].id, floor(count.index / 2))
  vpc_security_group_ids = [aws_security_group.tidb_sg.id]
  key_name               = "tidb-demo"

  root_block_device {
    volume_type = "gp3"
    volume_size = 20
  }

  tags = {
    Name = "TIDB-${count.index}"
  }
}

resource "aws_instance" "tikv" {
  count         = 3 * var.num_clusters
  ami           = var.ami_id
  instance_type = var.tikv_instance_type

  subnet_id              = element(aws_subnet.private_subnet[*].id, floor(count.index / 3))
  vpc_security_group_ids = [aws_security_group.tidb_sg.id]
  key_name               = "tidb-demo"

  root_block_device {
    volume_type = "gp3"
    volume_size = 500
  }

  tags = {
    Name = "TIKV-${count.index}"
  }
}

resource "aws_instance" "pd" {
  count         = 3 * var.num_clusters
  ami           = var.ami_id
  instance_type = var.pd_instance_type

  subnet_id              = element(aws_subnet.private_subnet[*].id, floor(count.index / 3))
  vpc_security_group_ids = [aws_security_group.tidb_sg.id]
  key_name               = "tidb-demo"

  root_block_device {
    volume_type = "gp3"
    volume_size = 20
  }

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

resource "aws_instance" "tiflash" {
  count         = var.num_clusters
  ami           = var.ami_id
  instance_type = var.tiflash_instance_type

  subnet_id              = aws_subnet.private_subnet[count.index].id
  vpc_security_group_ids = [aws_security_group.tidb_sg.id]
  key_name               = "tidb-demo"

  root_block_device {
    volume_type = "gp3"
    volume_size = 500
  }

  tags = {
    Name = "TiFlash-${count.index}"
  }
}

output "bastion_public_ips" {
  value = aws_instance.bastion[*].public_ip
}

output "pd_private_ips" {
  value = aws_instance.pd[*].private_ip
}

output "tidb_private_ips" {
  value = aws_instance.tidb[*].private_ip
}

output "tikv_private_ips" {
  value = aws_instance.tikv[*].private_ip
}

output "tiflash_private_ips" {
  value = aws_instance.tiflash[*].private_ip
}

############################################
# Generate CSV
############################################

locals {
  # Generate CSV lines for each server type
  all_csv_lines = concat(
    # Bastion host entry
    [for idx in range(length(aws_instance.bastion)) :
      "Demo,${aws_instance.bastion[idx].tags.Name},Bastion,${aws_instance.bastion[idx].public_ip},${aws_instance.bastion[idx].private_ip}"
    ],
    # PD server entries
    [for idx in range(length(aws_instance.pd)) :
      "Demo,${aws_instance.pd[idx].tags.Name},PD,${aws_instance.pd[idx].public_ip},${aws_instance.pd[idx].private_ip}"
    ],
    # TiKV server entries
    [for idx in range(length(aws_instance.tikv)) :
      "Demo,${aws_instance.tikv[idx].tags.Name},TiKV,${aws_instance.tikv[idx].public_ip},${aws_instance.tikv[idx].private_ip}"
    ],
    # TiDB server entries
    [for idx in range(length(aws_instance.tidb)) :
      "Demo,${aws_instance.tidb[idx].tags.Name},TiDB,${aws_instance.tidb[idx].public_ip},${aws_instance.tidb[idx].private_ip}"
    ],
    # TiFlash server entries
    [for idx in range(length(aws_instance.tiflash)) :
      "Demo,${aws_instance.tiflash[idx].tags.Name},TiFlash,${aws_instance.tiflash[idx].public_ip},${aws_instance.tiflash[idx].private_ip}"
    ],
    # Monitoring, Grafana, and Alertmanager roles on PD0
    [for idx in range(length(aws_instance.pd)) :
      "Demo,${aws_instance.pd[idx].tags.Name},Monitoring,Grafana,Alertmanager,${aws_instance.pd[idx].public_ip},${aws_instance.pd[idx].private_ip}"
    ]
  )
}

resource "null_resource" "generate_csv" {
  provisioner "local-exec" {
    command = <<-EOT
      #!/bin/bash
      echo "Environment,Node Name,Role,Public IP,Private IP" > cluster_inventory.csv
      echo "${join("\n", local.all_csv_lines)}" >> cluster_inventory.csv
    EOT
  }

  depends_on = [
    aws_instance.bastion,
    aws_instance.pd,
    aws_instance.tikv,
    aws_instance.tidb,
    aws_instance.tiflash
  ]
}