# main.tf
provider "aws" {
  region = var.region
}

terraform {
  backend "s3" {
  }
}

locals {
  vpc_lattice_region_unsupported_azs = {
    # US East (N. Virginia)
    "us-east-1" = ["use1-az3"]

    # US East (Ohio)
    "us-east-2" = []

    # US West (N. California)
    "us-west-1" = ["usw1-az2"]

    # US West (Oregon)
    "us-west-2" = []

    # Africa (Cape Town)
    "af-south-1" = []

    # Asia Pacific (Mumbai)
    "ap-south-1" = []

    # Asia Pacific (Singapore)
    "ap-southeast-1" = []

    # Asia Pacific (Sydney)
    "ap-southeast-2" = []

    # Asia Pacific (Tokyo)
    "ap-northeast-1" = ["apne1-az3"]

    # Asia Pacific (Seoul)
    "ap-northeast-2" = ["apne2-az2"]

    # Canada (Central)
    "ca-central-1" = ["cac1-az3"]

    # Europe (Frankfurt)
    "eu-central-1" = ["euc1-az2"]

    # Europe (Ireland)
    "eu-west-1" = ["euw1-az4"]

    # Europe (London)
    "eu-west-2" = []

    # Europe (Paris)
    "eu-west-3" = []

    # Europe (Milan)
    "eu-south-1" = []

    # Europe (Stockholm)
    "eu-north-1" = []

    # South America (São Paulo)
    "sa-east-1" = []

    # Israel (Tel Aviv)
    "il-central-1" = ["ilc1-az2"]
  }

  rds_proxy_unsupported_azs = {
    # US East (N. Virginia)
    "us-east-1" = ["use1-az3"]

    # US West (N. California)
    "us-west-1" = ["usw1-az2"]
  }
}

locals {
  is_supported_vpc_lattice_region = contains(
    keys(local.vpc_lattice_region_unsupported_azs),
    var.region
  )

  use_vpc_resource_gateway = !var.publicly_accessible && (
    var.privatelink_type == "auto" ? local.is_supported_vpc_lattice_region :
    var.privatelink_type == "resource-gateway" ? true :
    false
  )

  use_rds_proxy = !var.publicly_accessible && (
    var.privatelink_type == "auto" ? !local.is_supported_vpc_lattice_region :
    var.privatelink_type == "rds-proxy" ? true :
    false
  )
}


data "aws_availability_zones" "available" {
  state            = "available"
  exclude_zone_ids = local.use_vpc_resource_gateway ? lookup(local.vpc_lattice_region_unsupported_azs, var.region, []) : lookup(local.rds_proxy_unsupported_azs, var.region, [])
  filter {
    name   = "zone-type"
    values = ["availability-zone"]
  }
  filter {
    name   = "region-name"
    values = [var.region]
  }
  dynamic "filter" {
    for_each = var.vpc.availability_zones == null ? [] : [1]
    content {
      name   = "zone-id"
      values = split(",", var.vpc.availability_zones)
    }
  }
}

locals {
  max_azs         = min(3, length(data.aws_availability_zones.available.names))
  azs             = slice(data.aws_availability_zones.available.names, 0, local.max_azs)
  az_count        = length(local.azs)
  subnets         = cidrsubnets(var.vpc.cidr, 3, 3, 3, 3, 3, 3, 3, 3) # 每个 subnet 是 /19（/16 + 3）
  public_subnets  = slice(local.subnets, 0, local.az_count)
  private_subnets = slice(local.subnets, 4, 4 + local.az_count)

  tags = merge(
    var.tags,
    {
      rds         = var.name,
      environment = var.env,
    }
  )
}

resource "null_resource" "privatelink_type_check" {
  count = !var.publicly_accessible ? 1 : 0

  lifecycle {
    precondition {
      condition     = !(var.privatelink_type == "resource-gateway" && !local.is_supported_vpc_lattice_region)
      error_message = "privatelink_type=resource-gateway is not supported in region ${var.region}"
    }
  }
}

######################################
# 1. Secret
######################################
resource "random_password" "db_password" {
  length           = 20
  special          = true
  upper            = true
  numeric          = true
  override_special = "~!#$^*()-_+="
}

resource "aws_secretsmanager_secret" "db_password" {
  count                   = var.managed_secret ? 0 : 1
  name_prefix             = "${var.name}-password"
  recovery_window_in_days = 0
}

resource "aws_secretsmanager_secret_version" "db_password_value" {
  count     = var.managed_secret ? 0 : 1
  secret_id = aws_secretsmanager_secret.db_password[0].id
  secret_string = jsonencode({
    username = var.db.admin_user
    password = random_password.db_password.result
  })
}

######################################
# 2. VPC and RDS
######################################
module "vpc" {
  source  = "terraform-aws-modules/vpc/aws"
  version = "~> 5.0"

  name = "${var.name}-vpc"
  cidr = var.vpc.cidr
  azs  = local.azs

  enable_nat_gateway = var.publicly_accessible
  single_nat_gateway = true
  private_subnets    = local.private_subnets
  public_subnets     = local.public_subnets

  enable_dns_support         = true
  enable_dns_hostnames       = true
  manage_default_network_acl = false

  tags = local.tags
}

resource "aws_db_subnet_group" "rds" {
  name       = "${var.name}-subnet-group"
  subnet_ids = module.vpc.public_subnets

  tags = local.tags
}

resource "aws_db_subnet_group" "rds_private" {
  name       = "${var.name}-subnet-group-private"
  subnet_ids = module.vpc.private_subnets

  tags = local.tags
}

resource "aws_security_group" "rds_sg" {
  name        = "${var.name}-sg"
  description = "Allow DB access"
  vpc_id      = module.vpc.vpc_id

  dynamic "ingress" {
    for_each = !var.publicly_accessible || length(var.public_access_cidrs) > 0 ? [1] : []
    content {
      from_port   = var.db.port
      to_port     = var.db.port
      protocol    = "tcp"
      cidr_blocks = var.publicly_accessible ? var.public_access_cidrs : [module.vpc.vpc_cidr_block]
    }
  }

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

  tags = merge(
    local.tags,
    {
      Name = "${var.name}-sg"
    }
  )
}


resource "aws_rds_cluster_parameter_group" "this" {
  name        = var.name
  family      = "${var.db.engine}${var.db.engine_version}"
  description = "RDS Cluster Parameter Group for ${var.name}"

  dynamic "parameter" {
    for_each = var.cluster_params
    content {
      name  = parameter.key
      value = parameter.value
    }
  }

  tags = local.tags
}

resource "aws_rds_cluster" "managed_password" {
  count = var.managed_secret ? 1 : 0

  cluster_identifier              = var.name
  engine                          = var.db.engine
  engine_version                  = var.db.engine_version
  port                            = var.db.port
  database_name                   = var.db.default_db_name
  master_username                 = var.db.admin_user
  storage_type                    = var.db.storage_type
  allocated_storage               = var.db.allocated_storage
  storage_encrypted               = var.db.storage_encrypted
  skip_final_snapshot             = true
  db_subnet_group_name            = var.publicly_accessible ? aws_db_subnet_group.rds.name : aws_db_subnet_group.rds_private.name
  db_cluster_parameter_group_name = aws_rds_cluster_parameter_group.this.name
  vpc_security_group_ids          = [aws_security_group.rds_sg.id]
  deletion_protection             = true

  manage_master_user_password = true

  tags = merge(var.tags,
    {
      environment = var.env,
      Name        = var.name
    }
  )

  lifecycle {
    ignore_changes = [
      cluster_members
    ]
  }
}

resource "aws_rds_cluster" "self_managed_password" {
  count = var.managed_secret ? 0 : 1

  cluster_identifier              = var.name
  engine                          = var.db.engine
  engine_version                  = var.db.engine_version
  port                            = var.db.port
  database_name                   = var.db.default_db_name
  master_username                 = var.db.admin_user
  storage_type                    = var.db.storage_type
  allocated_storage               = var.db.allocated_storage
  storage_encrypted               = var.db.storage_encrypted
  skip_final_snapshot             = true
  db_subnet_group_name            = aws_db_subnet_group.rds.name
  db_cluster_parameter_group_name = aws_rds_cluster_parameter_group.this.name
  vpc_security_group_ids          = [aws_security_group.rds_sg.id]
  deletion_protection             = true

  master_password = random_password.db_password.result

  tags = merge(var.tags,
    {
      environment = var.env,
      Name        = var.name
    }
  )

  lifecycle {
    ignore_changes = [
      cluster_members
    ]
  }
}

locals {
  db_instance          = var.managed_secret ? aws_rds_cluster.managed_password[0] : aws_rds_cluster.self_managed_password[0]
  password_secret_arn  = var.managed_secret ? local.db_instance.master_user_secret[0].secret_arn : aws_secretsmanager_secret.db_password[0].arn
  password_secret_name = var.managed_secret ? local.db_instance.master_user_secret[0].secret_name : aws_secretsmanager_secret.db_password[0].name
}

resource "aws_rds_cluster_instance" "cluster_instances" {
  count               = var.db.instance_count
  identifier          = "${var.name}-${count.index}"
  cluster_identifier  = local.db_instance.id
  instance_class      = var.db.instance_class
  engine              = local.db_instance.engine
  engine_version      = local.db_instance.engine_version
  publicly_accessible = var.publicly_accessible

  tags = merge(
    local.tags,
    {
      Name = "${var.name}-instance-${count.index}"
    }
  )
}

######################################
# 3. VPC ResouceGateway
######################################
resource "aws_vpclattice_resource_gateway" "this" {
  count = local.use_vpc_resource_gateway ? 1 : 0

  name               = "${var.name}-rgw"
  vpc_id             = module.vpc.vpc_id
  subnet_ids         = module.vpc.private_subnets
  ip_address_type    = "IPV4"
  security_group_ids = [aws_security_group.rds_sg.id]

  tags = local.tags
}

resource "aws_vpclattice_resource_configuration" "this" {
  count = local.use_vpc_resource_gateway ? 1 : 0

  name = "${var.name}-rcfg"
  type = "ARN"

  resource_gateway_identifier                    = aws_vpclattice_resource_gateway.this[0].id
  allow_association_to_shareable_service_network = true

  resource_configuration_definition {
    arn_resource {
      arn = local.db_instance.arn
    }
  }

  tags = local.tags
}

######################################
# 4. RDS Proxy
######################################

resource "aws_iam_role" "rds_proxy_role" {
  count = local.use_rds_proxy ? 1 : 0

  name = "${var.name}-rds-proxy-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect = "Allow"
      Principal = {
        Service = "rds.amazonaws.com"
      }
      Action = "sts:AssumeRole"
    }]
  })
}

resource "aws_iam_role_policy_attachment" "rds_proxy_secrets" {
  count      = local.use_rds_proxy ? 1 : 0
  role       = aws_iam_role.rds_proxy_role[0].name
  policy_arn = "arn:aws:iam::aws:policy/SecretsManagerReadWrite"
}

resource "aws_db_proxy" "this" {
  count = local.use_rds_proxy ? 1 : 0

  name                   = "${var.name}-proxy"
  engine_family          = var.db.engine == "aurora-mysql" ? "MYSQL" : "POSTGRESQL"
  role_arn               = aws_iam_role.rds_proxy_role[0].arn
  vpc_subnet_ids         = module.vpc.private_subnets
  vpc_security_group_ids = [aws_security_group.rds_sg.id]
  require_tls            = false

  auth {
    auth_scheme               = "SECRETS"
    secret_arn                = local.password_secret_arn
    iam_auth                  = "DISABLED"
    client_password_auth_type = "MYSQL_NATIVE_PASSWORD" # MYSQL_CACHING_SHA2_PASSWORD requires TLS enabled
    description               = "rds master password"
  }

  tags = local.tags
}

resource "aws_db_proxy_default_target_group" "this" {
  count         = local.use_rds_proxy ? 1 : 0
  db_proxy_name = aws_db_proxy.this[0].name
}

resource "aws_db_proxy_target" "this" {
  count = local.use_rds_proxy ? 1 : 0

  db_proxy_name         = aws_db_proxy.this[0].name
  target_group_name     = aws_db_proxy_default_target_group.this[0].name
  db_cluster_identifier = local.db_instance.id
}

######################################
# 5. RDS Proxy NLB
######################################
resource "aws_lb" "rds_proxy_nlb" {
  count = local.use_rds_proxy ? 1 : 0

  name               = "${var.name}-proxy-nlb"
  load_balancer_type = "network"
  internal           = true
  subnets            = module.vpc.private_subnets

  tags = local.tags
}

resource "aws_lb_target_group" "rds_proxy_tg" {
  count = local.use_rds_proxy ? 1 : 0

  name        = "${var.name}-proxy-tg"
  port        = var.db.port
  protocol    = "TCP"
  vpc_id      = module.vpc.vpc_id
  target_type = "ip"

  health_check {
    port     = var.db.port
    protocol = "TCP"
  }

  tags = local.tags
}

resource "aws_lb_listener" "rds_proxy_listener" {
  count = local.use_rds_proxy ? 1 : 0

  load_balancer_arn = aws_lb.rds_proxy_nlb[0].arn
  port              = var.db.port
  protocol          = "TCP"

  default_action {
    type             = "forward"
    target_group_arn = aws_lb_target_group.rds_proxy_tg[0].arn
  }
  tags = merge(
    local.tags,
    {
      Name = "${var.name}-proxy-listener"
    }
  )
}

# data "external" "rds_proxy_ip" {
#   count   = local.use_rds_proxy ? 1 : 0
#   program = ["bash", "${path.module}/resolve-ip.sh"]

#   query = {
#     dns_name = aws_db_proxy.this[0].endpoint
#   }
# }

# resource "aws_lb_target_group_attachment" "rds_proxy_attachments" {
#   count            = local.use_rds_proxy ? 1 : 0
#   target_group_arn = aws_lb_target_group.rds_proxy_tg[0].arn
#   target_id        = data.external.rds_proxy_ip[0].result["ip"]
#   port             = var.db.port

#   lifecycle {
#     ignore_changes = [target_id]
#   }
# }

######################################
# 6. RDS Proxy Endpoint Service
######################################
resource "aws_vpc_endpoint_service" "rds_proxy" {
  count = local.use_rds_proxy ? 1 : 0

  acceptance_required        = false
  network_load_balancer_arns = [aws_lb.rds_proxy_nlb[0].arn]

  tags = merge(
    local.tags,
    {
      Name = "${var.name}-rds-proxy-service"
    }
  )
}

resource "null_resource" "wait_for_endpoint_service" {
  count = local.use_rds_proxy ? 1 : 0

  provisioner "local-exec" {
    interpreter = ["/bin/bash", "-c"]
    command     = <<EOT
      SERVICE_ID=${aws_vpc_endpoint_service.rds_proxy[0].id}
      MAX_ATTEMPTS=30
      SLEEP=10

      for ((i=0; i<MAX_ATTEMPTS; i++)); do
        state=$(aws ec2 describe-vpc-endpoint-service-configurations \
          --service-ids $SERVICE_ID \
          --query 'ServiceConfigurations[0].ServiceState' \
          --region ${var.region} \
          --output text 2>/dev/null || echo "error")

        echo "Attempt $((i+1))/$MAX_ATTEMPTS - Service state: $state"

        if [ "$${state,,}" == "available" ]; then
          echo "✅ Service is available!"
          exit 0
        elif [ "$state" == "error" ]; then
          echo "❌ Error checking service state"
          exit 1
        fi

        sleep $SLEEP
      done

      echo "⏰ Timeout reached after $((MAX_ATTEMPTS*SLEEP)) seconds"
      exit 1
    EOT
  }

  depends_on = [
    aws_vpc_endpoint_service.rds_proxy[0],
    local.db_instance
  ]
}


######################################
# 7. Lambda to refresh RDS Proxy target
######################################
resource "aws_iam_role" "lambda_rds_proxy" {
  count = local.use_rds_proxy ? 1 : 0

  name = "${var.name}-lambda-rds-proxy-updater"
  assume_role_policy = jsonencode({
    Version = "2012-10-17",
    Statement = [
      {
        Action = "sts:AssumeRole",
        Effect = "Allow",
        Principal = {
          Service = "lambda.amazonaws.com"
        }
      }
    ]
  })
  tags = local.tags
}

resource "aws_iam_role_policy" "lambda_rds_proxy" {
  count = local.use_rds_proxy ? 1 : 0

  name = "${var.name}-lambda-rds-proxy-updater-policy"
  role = aws_iam_role.lambda_rds_proxy[0].id
  policy = jsonencode({
    Version = "2012-10-17",
    Statement = [
      {
        Effect = "Allow",
        Action = [
          "elasticloadbalancing:RegisterTargets",
          "elasticloadbalancing:DeregisterTargets",
          "elasticloadbalancing:DescribeTargetHealth",
          "logs:CreateLogGroup",
          "logs:CreateLogStream",
          "logs:PutLogEvents"
        ],
        Resource = "*"
      }
    ]
  })
}

data "archive_file" "lambda_rds_proxy_updater" {
  count = local.use_rds_proxy ? 1 : 0

  type        = "zip"
  source_file = "${path.module}/lambda_rds_proxy_updater.py"
  output_path = "${path.module}/lambda_rds_proxy_updater.zip"
}

resource "aws_lambda_function" "rds_proxy_updater" {
  count = local.use_rds_proxy ? 1 : 0

  function_name = "${var.name}-rds-proxy-nlb-updater"
  role          = aws_iam_role.lambda_rds_proxy[0].arn
  handler       = "lambda_rds_proxy_updater.lambda_handler"
  runtime       = "python3.12"

  filename         = data.archive_file.lambda_rds_proxy_updater[0].output_path
  source_code_hash = data.archive_file.lambda_rds_proxy_updater[0].output_base64sha256

  environment {
    variables = {
      RDS_PROXY_ENDPOINT = aws_db_proxy.this[0].endpoint
      TARGET_GROUP_ARN   = aws_lb_target_group.rds_proxy_tg[0].arn
      TARGET_PORT        = tostring(var.db.port)
    }
  }
  tags = local.tags
}

resource "null_resource" "invoke_updater_once" {
  count = local.use_rds_proxy ? 1 : 0

  triggers = {
    lambda_version = aws_lambda_function.rds_proxy_updater[0].version
  }

  provisioner "local-exec" {
    command = <<EOT
      payload=$(echo '{"action":"init"}' | base64)
      aws lambda invoke \
        --function-name ${aws_lambda_function.rds_proxy_updater[0].function_name} \
        --payload $payload \
        --region ${var.region} \
        response.json
    EOT
  }
}

resource "aws_cloudwatch_event_rule" "every_5_min" {
  count = local.use_rds_proxy ? 1 : 0

  name                = "rds-proxy-updater-schedule"
  schedule_expression = "rate(5 minutes)"
  tags                = local.tags
}

resource "aws_cloudwatch_event_target" "lambda_target" {
  count = local.use_rds_proxy ? 1 : 0

  rule      = aws_cloudwatch_event_rule.every_5_min[0].name
  target_id = "${var.name}-RdsProxyUpdaterLambda"
  arn       = aws_lambda_function.rds_proxy_updater[0].arn
}

resource "aws_lambda_permission" "allow_eventbridge" {
  count = local.use_rds_proxy ? 1 : 0

  statement_id  = "AllowExecutionFromEventBridge"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.rds_proxy_updater[0].function_name
  principal     = "events.amazonaws.com"
  source_arn    = aws_cloudwatch_event_rule.every_5_min[0].arn
}