{ lib, pkgs, ... }:
let
  serverSize = "s-2vcpu-2gb";

  extraFilesScript = pkgs.writeShellScript "extra-files-script" ''
    #!/usr/bin/env bash
    set -euo pipefail

    mkdir -p etc/ssh/

    if [ -n "''${DO_SSH_HOSTKEY:-}" ]; then
      echo "Setting up SSH host key from environment"
      echo "$DO_SSH_HOSTKEY" | base64 -d > etc/ssh/ssh_host_ed25519_key
      chmod 0600 etc/ssh/ssh_host_ed25519_key
    else
      echo "Warning: DO_SSH_HOSTKEY environment variable not set"
    fi
  '';

in
{
  provider.digitalocean = {
    # Token will be read from DIGITALOCEAN_TOKEN environment variable
  };

  resource = {
    # Random string for unique naming
    random_string.host = {
      length = 6;
      special = false;
      upper = false;
    };

    digitalocean_ssh_key.default = {
      name = "nixos-anywhere-\${random_string.host.result}";
      public_key = lib.tfRef "var.digitalocean_public_key";
    };

    digitalocean_droplet.nixos = {
      name = "nixos-\${random_string.host.result}";
      image = "ubuntu-24-04-x64"; # Bootstrap image
      size = serverSize;
      region = lib.tfRef "var.digitalocean_region";
      ssh_keys = [ "\${digitalocean_ssh_key.default.id}" ];
      tags = [
        "nixos"
        "infrastructure"
      ];
    };
  };

  module = {
    nixos-system-build = {
      source = "github.com/nix-community/nixos-anywhere//terraform/nix-build";
      attribute = ".#nixosConfigurations.do-rproxy.config.system.build.toplevel";
    };

    nixos-disko = {
      source = "github.com/nix-community/nixos-anywhere//terraform/nix-build";
      attribute = ".#nixosConfigurations.do-rproxy.config.system.build.diskoScript";
    };

    nixos-install = {
      source = "github.com/nix-community/nixos-anywhere//terraform/install";
      nixos_system = "\${module.nixos-system-build.result.out}";
      nixos_partitioner = "\${module.nixos-disko.result.out}";
      target_host = "\${digitalocean_droplet.nixos.ipv4_address}";
      build_on_remote = true;
      extra_files_script = toString extraFilesScript;
    };
  };

  output = {
    server_ip = {
      description = "IP address of the NixOS server";
      value = "\${digitalocean_droplet.nixos.ipv4_address}";
    };

    ssh_command = {
      description = "SSH command to connect to the server";
      value = "ssh root@\${digitalocean_droplet.nixos.ipv4_address}";
    };

    server_name = {
      description = "Name of the created server";
      value = "\${digitalocean_droplet.nixos.name}";
    };
  };
}