# ── Vault LB — GCP Global HTTPS LB ───────────────────────────────────────────
#
# LB frontend is unchanged: forwarding rule → target HTTPS proxy → url map → backend.
# Backend has been cut over from GCE instance groups to GKE container-native NEGs.
#
# Architecture (GKE):
#   - Vault Helm chart runs 3 pods in namespace: vault on GKE neuron-platform cluster
#   - GKE auto-creates zonal container-native NEGs via Service annotation:
#       cloud.google.com/neg: '{"exposed_ports":{"8200":{}}}'
#   - NEG backends replace GCE instance group backends in the Global HTTPS LB
#   - DNS unchanged: vault.neuralplatform.ai → 34.54.164.21 (same GCP Global LB IP)
#
# GCE nodes (vault-node-1/2/3) and their instance groups are retained in this file
# but no longer serve traffic. Remove them after validating GKE Vault is stable.
#
# NEG names: confirm with:
#   gcloud compute network-endpoint-groups list \
#     --filter="name~k8s1-f08a4c22" --project neuron-785695 --format="table(name,zone)"
# Then fill in local.vault_neg_names below and run terraform apply.

locals {
  vault_version = "1.19.2"

  vault_nodes = {
    "vault-node-1" = { zone = "us-central1-a", node_id = "vault-node-1" }
    "vault-node-2" = { zone = "us-central1-b", node_id = "vault-node-2" }
    "vault-node-3" = { zone = "us-central1-c", node_id = "vault-node-3" }
  }

  # Container-native NEG names created by GKE from the Vault Service annotation:
  #   cloud.google.com/neg: '{"exposed_ports":{"8200":{}}}'
  # GKE creates one NEG per zone where Vault pods are scheduled.
  # Confirmed with: gcloud compute network-endpoint-groups list --project neuron-785695
  vault_neg_names = {
    "us-central1-b" = "k8s1-bfbeff02-vault-vault-helm-gke-8200-db1a474f"
    "us-central1-c" = "k8s1-bfbeff02-vault-vault-helm-gke-8200-db1a474f"
    "us-central1-f" = "k8s1-bfbeff02-vault-vault-helm-gke-8200-db1a474f"
  }
}

# ── GKE container-native NEG data sources ────────────────────────────────────
# GKE auto-creates these when the Vault Service annotation is applied.
# Populate local.vault_neg_names above with the actual NEG names, then apply.

data "google_compute_network_endpoint_group" "vault_gke" {
  for_each = local.vault_neg_names
  name     = each.value
  zone     = each.key
  project  = var.project_id
}

# ── Service Account for Vault nodes ──────────────────────────────────────────
# Reuses the vault-unseal SA from vault-kms.tf for KMS access.
# Additional roles: logging + metrics from ops agent.

resource "google_project_iam_member" "vault_node_log_writer" {
  project = var.project_id
  role    = "roles/logging.logWriter"
  member  = "serviceAccount:${google_service_account.vault_unseal.email}"
}

resource "google_project_iam_member" "vault_node_metric_writer" {
  project = var.project_id
  role    = "roles/monitoring.metricWriter"
  member  = "serviceAccount:${google_service_account.vault_unseal.email}"
}

# Allow Vault nodes to read their own instance metadata (needed for GCP auth method later)
resource "google_project_iam_member" "vault_node_compute_viewer" {
  project = var.project_id
  role    = "roles/compute.viewer"
  member  = "serviceAccount:${google_service_account.vault_unseal.email}"
}

# ── Startup script staged in GCS ─────────────────────────────────────────────
# Stored in the existing runner-assets bucket (reuse infrastructure).
# The script installs Vault, writes the config, and starts the systemd unit.

resource "google_storage_bucket_object" "vault_startup" {
  name   = "vault/startup.sh"
  bucket = google_storage_bucket.runner_assets.name
  source = "${path.module}/vault/startup.sh"

  metadata = {
    sha256 = filesha256("${path.module}/vault/startup.sh")
  }
}

resource "google_storage_bucket_iam_member" "vault_node_bucket_read" {
  bucket = google_storage_bucket.runner_assets.name
  role   = "roles/storage.objectViewer"
  member = "serviceAccount:${google_service_account.vault_unseal.email}"
}

# ── Vault node VMs ────────────────────────────────────────────────────────────

resource "google_compute_instance" "vault_node" {
  for_each     = local.vault_nodes
  name         = each.key
  machine_type = "e2-small"
  zone         = each.value.zone
  project      = var.project_id

  tags = ["vault-node", "allow-iap-ssh"]

  boot_disk {
    initialize_params {
      image = "projects/debian-cloud/global/images/family/debian-12"
      size  = 20
      type  = "pd-balanced"
    }
  }

  # Separate persistent disk for Raft data — survives VM recreation
  attached_disk {
    source      = google_compute_disk.vault_data[each.key].self_link
    device_name = "vault-data"
    mode        = "READ_WRITE"
  }

  network_interface {
    network = "default"
    # No external IP — accessed via IAP SSH and the internal LB
    # Vault API is published externally via the regional LB below
    access_config {}
  }

  service_account {
    email  = google_service_account.vault_unseal.email
    scopes = ["cloud-platform"]
  }

  metadata = {
    # Pull and execute the real startup script from GCS
    startup-script = <<-EOT
      #!/usr/bin/env bash
      set -euxo pipefail
      apt-get update -y
      apt-get install -y curl ca-certificates apt-transport-https gnupg google-cloud-cli
      gsutil cat gs://${google_storage_bucket.runner_assets.name}/${google_storage_bucket_object.vault_startup.name} \
        > /tmp/vault-startup.sh
      chmod +x /tmp/vault-startup.sh
      VAULT_NODE_ID="${each.value.node_id}" \
      VAULT_VERSION="${local.vault_version}" \
        /tmp/vault-startup.sh
    EOT

    enable-oslogin = "TRUE"
  }

  allow_stopping_for_update = true

  depends_on = [
    google_compute_disk.vault_data,
    google_storage_bucket_object.vault_startup,
    google_storage_bucket_iam_member.vault_node_bucket_read,
  ]
}

# ── Persistent data disks ─────────────────────────────────────────────────────
# 10 GiB per node. Kept separate from the boot disk so Raft data
# survives a full VM deletion and recreation.

resource "google_compute_disk" "vault_data" {
  for_each = local.vault_nodes
  name     = "${each.key}-data"
  zone     = each.value.zone
  project  = var.project_id
  type     = "pd-ssd"
  size     = 10

  labels = {
    managed-by = "terraform"
    service    = "vault"
    node       = each.key
  }

  lifecycle {
    prevent_destroy = true
  }
}

# ── Firewall rules ────────────────────────────────────────────────────────────

# IAP SSH — ops access without a public SSH port
resource "google_compute_firewall" "vault_iap_ssh" {
  name      = "vault-iap-ssh"
  network   = "default"
  project   = var.project_id
  direction = "INGRESS"
  priority  = 1000

  allow {
    protocol = "tcp"
    ports    = ["22"]
  }

  source_ranges = ["35.235.240.0/20"] # GCP IAP CIDR
  target_tags   = ["vault-node"]
}

# Raft cluster traffic — node-to-node port 8201 (internal only)
resource "google_compute_firewall" "vault_raft" {
  name      = "vault-raft-internal"
  network   = "default"
  project   = var.project_id
  direction = "INGRESS"
  priority  = 1000

  allow {
    protocol = "tcp"
    ports    = ["8201"]
  }

  source_tags = ["vault-node"]
  target_tags = ["vault-node"]
}

# Vault API — allow from the GCP health check ranges and the LB
resource "google_compute_firewall" "vault_api" {
  name      = "vault-api-from-lb"
  network   = "default"
  project   = var.project_id
  direction = "INGRESS"
  priority  = 1000

  allow {
    protocol = "tcp"
    ports    = ["8200"]
  }

  # GCP health check ranges (130.211.0.0/22, 35.191.0.0/16)
  # and RFC1918 for any internal service access
  source_ranges = ["130.211.0.0/22", "35.191.0.0/16", "10.0.0.0/8", "172.16.0.0/12", "192.168.0.0/16"]
  target_tags   = ["vault-node"]
}

# GKE Vault pods — allow GCP health check probers to reach pod IPs on port 8200.
# Container-native NEGs direct health checks to pod IPs (not node IPs).
# GKE Autopilot pod CIDR: 10.45.128.0/22 (from cluster ip_allocation_policy).
# Without this rule, GCP LB health checks fail → "no healthy upstream".
# No target_tags — applies to all instances/pods in the default network.
# Safe: GCP health check ranges (130.211.0.0/22, 35.191.0.0/16) are GCP-internal only.
resource "google_compute_firewall" "vault_api_gke" {
  name      = "vault-api-from-lb-gke"
  network   = "default"
  project   = var.project_id
  direction = "INGRESS"
  priority  = 1000

  allow {
    protocol = "tcp"
    ports    = ["8200"]
  }

  # GCP health check source ranges only — restricted port, low risk
  source_ranges = ["130.211.0.0/22", "35.191.0.0/16"]
}

# ── Instance Groups (one unmanaged group per zone for the LB) ─────────────────

resource "google_compute_instance_group" "vault" {
  for_each  = local.vault_nodes
  name      = "vault-${each.key}"
  zone      = each.value.zone
  project   = var.project_id
  instances = [google_compute_instance.vault_node[each.key].self_link]

  named_port {
    name = "vault-api"
    port = 8200
  }
}

# ── Regional HTTPS LB for vault.neuralplatform.ai ─────────────────────────────
# We use a global external HTTPS LB (same scheme as the prod marketing LB)
# so we can attach a Google-managed cert for vault.neuralplatform.ai.
#
# TLS is terminated at the LB. Vault listens on plain 8200 internally.
# The Cloudflare DNS A record for vault.neuralplatform.ai (neuralplatform zone)
# must point to vault_lb_ip output below — add it in Cloudflare dashboard
# or in the legion Terraform if you bring that zone under TF management.

resource "google_compute_global_address" "vault" {
  name    = "vault-ip"
  project = var.project_id
}

resource "google_compute_managed_ssl_certificate" "vault" {
  name    = "vault-cert"
  project = var.project_id

  managed {
    domains = ["vault.neuralplatform.ai"]
  }
}

resource "google_compute_health_check" "vault" {
  name    = "vault-health"
  project = var.project_id

  # Vault listens on plain HTTP (tls_disable = true) — TLS is terminated at the LB
  http_health_check {
    port         = 8200
    request_path = "/v1/sys/health?standbyok=true&sealedok=true&uninitcode=200"
  }

  check_interval_sec  = 15
  timeout_sec         = 5
  healthy_threshold   = 2
  unhealthy_threshold = 3
}

resource "google_compute_backend_service" "vault" {
  name                  = "vault-backend"
  project               = var.project_id
  load_balancing_scheme = "EXTERNAL_MANAGED"
  protocol              = "HTTP"   # Vault serves plain HTTP; TLS terminates at the LB
  timeout_sec           = 30

  # port_name not used with NEG backends (NEGs have explicit port in the service annotation)

  health_checks = [google_compute_health_check.vault.self_link]

  # GKE container-native NEG backends — one per zone where Vault pods are scheduled.
  # RATE balancing mode is required for NEGs with EXTERNAL_MANAGED load balancers.
  dynamic "backend" {
    for_each = data.google_compute_network_endpoint_group.vault_gke
    content {
      group                 = backend.value.self_link
      balancing_mode        = "RATE"
      max_rate_per_endpoint = 100
    }
  }

  log_config {
    enable      = true
    sample_rate = 1.0
  }
}

resource "google_compute_url_map" "vault" {
  name            = "vault-urlmap"
  project         = var.project_id
  default_service = google_compute_backend_service.vault.self_link
}

resource "google_compute_target_https_proxy" "vault" {
  name             = "vault-https-proxy"
  project          = var.project_id
  url_map          = google_compute_url_map.vault.self_link
  ssl_certificates = [google_compute_managed_ssl_certificate.vault.self_link]
}

resource "google_compute_target_http_proxy" "vault_redirect" {
  name    = "vault-http-proxy"
  project = var.project_id
  url_map = google_compute_url_map.vault_redirect.self_link
}

resource "google_compute_url_map" "vault_redirect" {
  name    = "vault-urlmap-http-redirect"
  project = var.project_id

  default_url_redirect {
    https_redirect         = true
    redirect_response_code = "MOVED_PERMANENTLY_DEFAULT"
    strip_query            = false
  }
}

resource "google_compute_global_forwarding_rule" "vault_https" {
  name                  = "vault-fwd-https"
  project               = var.project_id
  target                = google_compute_target_https_proxy.vault.self_link
  ip_address            = google_compute_global_address.vault.address
  port_range            = "443"
  load_balancing_scheme = "EXTERNAL_MANAGED"
}

resource "google_compute_global_forwarding_rule" "vault_http" {
  name                  = "vault-fwd-http"
  project               = var.project_id
  target                = google_compute_target_http_proxy.vault_redirect.self_link
  ip_address            = google_compute_global_address.vault.address
  port_range            = "80"
  load_balancing_scheme = "EXTERNAL_MANAGED"
}

# ── Outputs ───────────────────────────────────────────────────────────────────

output "vault_lb_ip" {
  description = "Global IP for vault.neuralplatform.ai — set as DNS A record in Cloudflare (neuralplatform.ai zone)"
  value       = google_compute_global_address.vault.address
}

output "vault_node_zones" {
  description = "Zone placement of each Vault node"
  value = {
    for k, v in local.vault_nodes : k => v.zone
  }
}