// Node bootstrap. Everything local to a node is applied by cloud-init at
// first boot; the cluster forms without any manual step. Slurm runs in
// configless mode: only the head node carries slurm.conf, compute nodes
// fetch it from the controller. The munge key is generated per deployment
// and delivered to every node through its user data.

locals {
  munge_key_b64 = random_bytes.munge_key.base64

  // Written early by cloud-init to a staging path (the munge user does not
  // exist yet); runcmd installs it with the right owner after the packages.
  munge_stage_path = "/run/simpl/munge.key"

  cluster_hosts = concat(
    ["${local.head_ip} ${local.head_hostname}"],
    [for i in range(var.compute_node_count) : "${local.compute_ips[i]} ${local.compute_hostnames[i]}"],
  )

  slurm_conf = join("\n", concat([
    "ClusterName=${lower(var.name_prefix)}-hpc",
    "SlurmctldHost=${local.head_hostname}(${local.head_ip})",
    "SlurmctldParameters=enable_configless",
    "AuthType=auth/munge",
    "SlurmUser=slurm",
    "CredType=cred/munge",
    "SelectType=select/cons_tres",
    "SelectTypeParameters=CR_Core",
    "ProctrackType=proctrack/linuxproc",
    "TaskPlugin=task/none",
    "ReturnToService=2",
    "SlurmctldPort=6817",
    "SlurmdPort=6818",
    "SrunPortRange=60001-60128",
    "StateSaveLocation=/var/lib/slurm/slurmctld",
    "SlurmdSpoolDir=/var/lib/slurm/slurmd",
    "SlurmctldLogFile=/var/log/slurm/slurmctld.log",
    "SlurmdLogFile=/var/log/slurm/slurmd.log",
    "MailProg=/bin/true",
    ], [
    for i in range(var.compute_node_count) :
    "NodeName=${local.compute_hostnames[i]} NodeAddr=${local.compute_ips[i]} CPUs=${var.compute_cpus} State=UNKNOWN"
    ], [
    "PartitionName=main Nodes=ALL Default=YES MaxTime=INFINITE State=UP",
  ]))

  hosts_runcmd = [for entry in local.cluster_hosts : "grep -qxF '${entry}' /etc/hosts || echo '${entry}' >> /etc/hosts"]

  munge_runcmd = [
    "install -o munge -g munge -m 0400 ${local.munge_stage_path} /etc/munge/munge.key",
    "rm -f ${local.munge_stage_path}",
    "systemctl restart munge",
  ]

  // Every cluster node is on the private subnet only: the default route
  // comes from DHCP (the subnet gateway) but no name servers are provided,
  // so resolution is pinned to public resolvers. Runs on every boot and
  // before package installation; the last entry blocks until egress through
  // the gateway answers, capped at ~4 minutes.
  egress_bootcmd = [
    "mkdir -p /etc/systemd/resolved.conf.d",
    "printf '[Resolve]\\nDNS=1.1.1.1 8.8.8.8\\n' > /etc/systemd/resolved.conf.d/90-cluster.conf",
    "systemctl try-restart systemd-resolved || true",
    "for i in $(seq 1 48); do timeout 3 bash -c '</dev/tcp/archive.ubuntu.com/80' 2>/dev/null && break; sleep 5; done",
  ]

  head_cloud_config = {
    hostname                  = local.head_hostname
    prefer_fqdn_over_hostname = false
    package_update            = true
    packages = [
      "slurmctld",
      "slurm-client",
      "munge",
      "nfs-kernel-server",
      "openmpi-bin",
      "chrony",
    ]
    write_files = [
      {
        path        = local.munge_stage_path
        encoding    = "b64"
        content     = local.munge_key_b64
        permissions = "0400"
      },
      {
        path        = "/etc/slurm/slurm.conf"
        content     = "${local.slurm_conf}\n"
        permissions = "0644"
      },
      {
        path        = "/etc/exports"
        content     = "/home ${var.cluster_cidr}(rw,sync,no_subtree_check,root_squash)\n"
        permissions = "0644"
      },
    ]
    bootcmd = local.egress_bootcmd
    runcmd = concat(local.hosts_runcmd, local.munge_runcmd, [
      "mkdir -p /var/lib/slurm/slurmctld /var/log/slurm",
      "chown slurm:slurm /var/lib/slurm/slurmctld /var/log/slurm",
      "systemctl enable --now nfs-kernel-server",
      "exportfs -ra",
      "systemctl enable slurmctld",
      "systemctl restart slurmctld",
    ])
  }

  compute_cloud_config = [for i in range(var.compute_node_count) : {
    hostname                  = local.compute_hostnames[i]
    prefer_fqdn_over_hostname = false
    package_update            = true
    packages = [
      "slurmd",
      "munge",
      "nfs-common",
      "openmpi-bin",
      "chrony",
    ]
    bootcmd = local.egress_bootcmd
    write_files = [
      {
        path        = local.munge_stage_path
        encoding    = "b64"
        content     = local.munge_key_b64
        permissions = "0400"
      },
      {
        path        = "/etc/default/slurmd"
        content     = "SLURMD_OPTIONS=\"--conf-server ${local.head_ip}:6817\"\n"
        permissions = "0644"
      },
      {
        path        = "/etc/systemd/system/slurmd.service.d/90-retry.conf"
        content     = "[Unit]\nStartLimitIntervalSec=0\n\n[Service]\nRestart=on-failure\nRestartSec=15\n"
        permissions = "0644"
      },
    ]
    mounts = [
      ["${local.head_ip}:/home", "/home", "nfs", "defaults,nofail,_netdev,x-systemd.automount", "0", "0"],
    ]
    runcmd = concat(local.hosts_runcmd, local.munge_runcmd, [
      "mkdir -p /var/lib/slurm/slurmd /var/log/slurm",
      "chown slurm:slurm /var/lib/slurm/slurmd /var/log/slurm",
      // Best effort early mount of the shared home, so jobs submitted right
      // after the cluster forms never land on the local directory before the
      // automount first triggers.
      "systemctl start home.mount 2>/dev/null || true",
      // The retry drop-in keeps slurmd re-launching until the controller
      // answers: the cluster nodes boot in parallel and the first attempts
      // may precede slurmctld.
      "systemctl daemon-reload",
      "systemctl enable slurmd",
      "systemctl restart slurmd || true",
    ])
  }]

  // Consumer and management users, shared by every node. Fixed uids: Slurm
  // identifies users by uid across the cluster, so the consumer must exist
  // everywhere with the same one. The consumer has no sudo (it is a cluster
  // user, not a machine owner) and its password exists only on the head node
  // (applied over SSH, never through user data). merge directives make
  // cloud-init append to the node lists instead of replacing them.
  credentials_cloud_config = {
    merge_how = [
      { name = "list", settings = ["append"] },
      { name = "dict", settings = ["no_replace", "recurse_list"] },
    ]
    ssh_pwauth = true
    users = [
      "default",
      {
        name        = var.username
        uid         = 2000
        shell       = "/bin/bash"
        lock_passwd = false
        groups      = "users"
      },
      {
        name                = var.admin_username
        uid                 = 2001
        shell               = "/bin/bash"
        lock_passwd         = true
        sudo                = "ALL=(ALL) NOPASSWD:ALL"
        groups              = "users, sudo"
        ssh_authorized_keys = [trimspace(tls_private_key.admin.public_key_openssh)]
      },
      // The munge and slurm system users are created before the packages so
      // they get the SAME uid on every node: the packages otherwise allocate
      // them dynamically, and Slurm authenticates controller-to-daemon
      // messages by uid ("Security violation, slurm message from uid N"
      // across nodes when they diverge). Low fixed uids: the dynamic system
      // range is allocated downward from 999, so the high 900s are already
      // taken on the stock image. Listed last so a failure here can never
      // block the login users above.
      {
        name           = "munge"
        uid            = 601
        system         = true
        shell          = "/usr/sbin/nologin"
        homedir        = "/var/lib/munge"
        no_create_home = true
        lock_passwd    = true
      },
      {
        name           = "slurm"
        uid            = 602
        system         = true
        shell          = "/usr/sbin/nologin"
        homedir        = "/var/lib/slurm"
        no_create_home = true
        lock_passwd    = true
      },
    ]
  }
}

data "cloudinit_config" "head" {
  gzip          = false
  base64_encode = true

  part {
    content_type = "text/cloud-config"
    content      = join("\n", ["#cloud-config", yamlencode(local.head_cloud_config)])
  }

  part {
    content_type = "text/cloud-config"
    content      = join("\n", ["#cloud-config", yamlencode(local.credentials_cloud_config)])
  }
}

data "cloudinit_config" "compute" {
  count         = var.compute_node_count
  gzip          = false
  base64_encode = true

  part {
    content_type = "text/cloud-config"
    content      = join("\n", ["#cloud-config", yamlencode(local.compute_cloud_config[count.index])])
  }

  part {
    content_type = "text/cloud-config"
    content      = join("\n", ["#cloud-config", yamlencode(local.credentials_cloud_config)])
  }
}