{ pkgs, serverPkg, clientPkg }:
let
  serverModule = import ../nixos/modules/lightscale-server.nix {
    defaultPackage = serverPkg;
  };
  clientModule = import ../nixos/modules/lightscale-client.nix {
    defaultPackage = clientPkg;
  };
in
{
  name = "lightscale-lab-simple-autoreg";
  nodes = {
    server = { ... }: {
      imports = [ serverModule ];
      networking.hostName = "server";
      networking.usePredictableInterfaceNames = false;
      virtualisation.vlans = [ 1 ];
      networking.interfaces.eth1.useDHCP = false;
      networking.interfaces.eth1.ipv4.addresses = [
        {
          address = "10.0.0.1";
          prefixLength = 24;
        }
      ];
      networking.firewall.enable = false;
      boot.kernelModules = [ "wireguard" ];

      services.lightscale-server = {
        enable = true;
        listen = "10.0.0.1:8080";
        stateFile = "/var/lib/lightscale-server/state.json";
        adminToken = "lab-admin-token";
      };

      environment.systemPackages = [
        clientPkg
        pkgs.curl
        pkgs.iputils
      ];
    };

    client1 = { ... }: {
      imports = [ clientModule ];
      networking.hostName = "client1";
      networking.usePredictableInterfaceNames = false;
      virtualisation.vlans = [ 1 ];
      networking.interfaces.eth1.useDHCP = false;
      networking.interfaces.eth1.ipv4.addresses = [
        {
          address = "10.0.0.2";
          prefixLength = 24;
        }
      ];
      networking.firewall.enable = false;
      boot.kernelModules = [ "wireguard" ];

      services.lightscale-client = {
        enable = true;
        controlUrls = [ "http://10.0.0.1:8080" ];
        autoRegister = true;
        enrollmentTokenFile = "/run/lightscale-enroll.token";
        registerNodeName = "simple-client1";
        registerExtraArgs = [
          "--machine-private-key-file"
          "/run/lightscale-machine.key"
          "--wg-private-key-file"
          "/run/lightscale-wg.key"
        ];
        agentArgs = [
          "--listen-port"
          "51820"
          "--apply-routes"
          "--heartbeat-interval"
          "5"
          "--longpoll-timeout"
          "5"
        ];
      };

      environment.systemPackages = [
        clientPkg
        pkgs.iputils
      ];
    };

    client2 = { ... }: {
      imports = [ clientModule ];
      networking.hostName = "client2";
      networking.usePredictableInterfaceNames = false;
      virtualisation.vlans = [ 1 ];
      networking.interfaces.eth1.useDHCP = false;
      networking.interfaces.eth1.ipv4.addresses = [
        {
          address = "10.0.0.3";
          prefixLength = 24;
        }
      ];
      networking.firewall.enable = false;
      boot.kernelModules = [ "wireguard" ];

      services.lightscale-client = {
        enable = true;
        controlUrls = [ "http://10.0.0.1:8080" ];
        autoRegister = true;
        enrollmentTokenFile = "/run/lightscale-enroll.token";
        registerNodeName = "simple-client2";
        registerExtraArgs = [
          "--machine-private-key-file"
          "/run/lightscale-machine.key"
          "--wg-private-key-file"
          "/run/lightscale-wg.key"
        ];
        agentArgs = [
          "--listen-port"
          "51820"
          "--apply-routes"
          "--heartbeat-interval"
          "5"
          "--longpoll-timeout"
          "5"
        ];
      };

      environment.systemPackages = [
        clientPkg
        pkgs.iputils
      ];
    };
  };

  testScript = ''
    import base64
    import json

    start_all()

    server.wait_for_unit("lightscale-server.service")
    server.wait_for_open_port(8080, addr="10.0.0.1", timeout=120)
    client1.wait_for_unit("multi-user.target")
    client2.wait_for_unit("multi-user.target")

    net = json.loads(server.succeed(
        "curl -sSf -X POST http://10.0.0.1:8080/v1/networks "
        "-H 'authorization: Bearer lab-admin-token' "
        "-H 'content-type: application/json' "
        "-d '{\"name\":\"simple-autoreg\",\"bootstrap_token_ttl_seconds\":1200,\"bootstrap_token_uses\":10}'"
    ))
    token = net["bootstrap_token"]["token"]
    network_id = net["network"]["id"]

    def write_registration_files(node, machine_byte, wg_byte):
        machine_key = base64.b64encode(bytes([machine_byte]) * 32).decode("ascii")
        wg_key = base64.b64encode(bytes([wg_byte]) * 32).decode("ascii")
        node.succeed(f"printf '%s\\n' '{token}' > /run/lightscale-enroll.token")
        node.succeed(f"printf '%s\\n' '{machine_key}' > /run/lightscale-machine.key")
        node.succeed(f"printf '%s\\n' '{wg_key}' > /run/lightscale-wg.key")
        return machine_key, wg_key

    def run_autoreg(node):
        node.execute("systemctl reset-failed lightscale-client-register.service || true")
        node.succeed("systemctl start lightscale-client-register.service")
        node.wait_for_unit("lightscale-client-register.service")
        node.wait_until_succeeds("test -s /var/lib/lightscale-client/default/state.json", timeout=120)
        node.wait_until_succeeds("systemctl is-active lightscale-client.service", timeout=120)
        node.wait_until_succeeds("ip link show ls-default", timeout=120)

    expected1 = write_registration_files(client1, 11, 21)
    expected2 = write_registration_files(client2, 12, 22)
    run_autoreg(client1)
    run_autoreg(client2)

    state1 = json.loads(client1.succeed("cat /var/lib/lightscale-client/default/state.json"))
    state2 = json.loads(client2.succeed("cat /var/lib/lightscale-client/default/state.json"))
    assert state1["network_id"] == network_id
    assert state2["network_id"] == network_id
    assert state1["machine_private_key"] == expected1[0]
    assert state1["wg_private_key"] == expected1[1]
    assert state2["machine_private_key"] == expected2[0]
    assert state2["wg_private_key"] == expected2[1]

    ip1 = state1["ipv4"]
    ip2 = state2["ipv4"]

    server.wait_until_succeeds(
        f"curl -sSf -H 'authorization: Bearer lab-admin-token' "
        f"http://10.0.0.1:8080/v1/admin/networks/{network_id}/nodes | grep -q '\"approved\":true'",
        timeout=120,
    )

    client1.wait_until_succeeds(f"ping -c 3 {ip2}", timeout=180)
    client2.wait_until_succeeds(f"ping -c 3 {ip1}", timeout=180)

    client1.succeed("systemctl restart lightscale-client.service")
    client1.wait_for_unit("lightscale-client.service")
    client1.wait_until_succeeds(f"ping -c 3 {ip2}", timeout=180)
  '';
}