/*
  There are two roles: peers and controllers:
    - Every controller has an endpoint set
    - There can be multiple peers
    - There has to be one or more controllers
    - Peers connect to ALL controllers (full mesh)
    - If only one controller exists, peers automatically use it for IP allocation
    - If multiple controllers exist, peers must specify which controller's subnet to use
    - Controllers have IPv6 forwarding enabled, every peer and controller can reach
      everyone else, via extra controller hops if necessary

    Example:
              ┌───────────────────────────────┐
              │            ◄─────────────     │
              │ controller2              controller1
              │    ▲       ─────────────►    ▲     ▲
              │    │ │ │ │                 │ │   │ │
              │    │ │ │ │                 │ │   │ │
              │    │ │ │ │                 │ │   │ │
              │    │ │ │ └───────────────┐ │ │   │ │
              │    │ │ └──────────────┐  │ │ │   │ │
              │      ▼                │  ▼ ▼     ▼
              └─► peer2               │  peer1  peer3
                                      │          ▲
                                      └──────────┘

  Network Architecture:

  IPv6 Address Allocation:
    - Base network: /40 ULA prefix (generated from instance name)
    - Controllers: Each gets a /56 subnet from the base /40
    - Peers: Each gets a unique host suffix that is used in ALL controller subnets

  Address Assignment:
    - Each peer generates a unique 64-bit host suffix (e.g., :8750:a09b:0:1)
    - This suffix is appended to each controller's /56 prefix
    - Example: peer1 with suffix :8750:a09b:0:1 gets:
      - fd51:19c1:3b:f700:8750:a09b:0:1 in controller1's subnet
      - fd51:19c1:c1:aa00:8750:a09b:0:1 in controller2's subnet

  Peers: Use a SINGLE interface that:
    - Connects to ALL controllers
    - Has multiple IPs, one in each controller's subnet (with /56 prefix)
    - Routes to each controller's /56 subnet via that controller
    - allowedIPs: Each controller's /56 subnet
    - No routing conflicts due to unique IPs per subnet

  Controllers: Use a SINGLE interface that:
    - Connects to ALL peers and ALL other controllers
    - Gets a /56 subnet from the base /40 network
    - Has IPv6 forwarding enabled for routing between peers
    - allowedIPs:
      - For peers: A /96 range containing the peer's address in this controller's subnet
      - For other controllers: The controller's /56 subnet
*/

{ ... }:
let
  # Shared module for extraHosts configuration
  extraHostsModule =
    {
      instanceName,
      settings,
      roles,
      config,
      lib,
      ...
    }:
    {
      networking.extraHosts =
        let
          domain = if settings.domain == null then instanceName else settings.domain;
          # Controllers use their subnet's ::1 address
          controllerHosts = lib.mapAttrsToList (
            name: _value:
            let
              prefix = builtins.readFile (
                config.clan.core.settings.directory
                + "/vars/per-machine/${name}/wireguard-network-${instanceName}/prefix/value"
              );
              # Controller IP is always ::1 in their subnet
              ip = prefix + "::1";
            in
            "${ip} ${name}.${domain}"
          ) roles.controller.machines;

          # Peers use their suffix in their designated controller's subnet only
          peerHosts = lib.mapAttrsToList (
            peerName: peerValue:
            let
              peerSuffix = builtins.readFile (
                config.clan.core.settings.directory
                + "/vars/per-machine/${peerName}/wireguard-network-${instanceName}/suffix/value"
              );
              # Determine designated controller
              designatedController =
                if (builtins.length (builtins.attrNames roles.controller.machines) == 1) then
                  (builtins.head (builtins.attrNames roles.controller.machines))
                else
                  peerValue.settings.controller;
              controllerPrefix = builtins.readFile (
                config.clan.core.settings.directory
                + "/vars/per-machine/${designatedController}/wireguard-network-${instanceName}/prefix/value"
              );
              peerIP = controllerPrefix + ":" + peerSuffix;
            in
            "${peerIP} ${peerName}.${domain}"
          ) roles.peer.machines;

          # External peers
          externalPeerHosts = lib.flatten (
            lib.mapAttrsToList (
              ctrlName: _ctrlValue:
              lib.map (
                peer:
                let
                  peerSuffix = builtins.readFile (
                    config.clan.core.settings.directory
                    + "/vars/shared/wireguard-network-${instanceName}-external-peer-${peer}/suffix/value"
                  );
                  controllerPrefix = builtins.readFile (
                    config.clan.core.settings.directory
                    + "/vars/per-machine/${ctrlName}/wireguard-network-${instanceName}/prefix/value"
                  );
                  peerIP = controllerPrefix + ":" + peerSuffix;
                in
                "${peerIP} ${peer}.${domain}"
              ) (roles.controller.machines.${ctrlName}.settings.externalPeers)
            ) roles.controller.machines
          );
        in
        builtins.concatStringsSep "\n" (controllerHosts ++ peerHosts ++ externalPeerHosts);
    };

  # Shared interface options
  sharedInterface =
    { lib, ... }:
    {
      options.port = lib.mkOption {
        type = lib.types.int;
        example = 51820;
        default = 51820;
        description = ''
          Port for the wireguard interface
        '';
      };

      options.domain = lib.mkOption {
        type = lib.types.nullOr lib.types.str;
        defaultText = lib.literalExpression "instanceName";
        default = null;
        description = ''
          Domain suffix to use for hostnames in /etc/hosts.
          Defaults to the instance name.
        '';
      };
    };
in
{
  _class = "clan.service";
  manifest.name = "clan-core/wireguard";
  manifest.description = "Wireguard-based VPN mesh network with automatic IPv6 address allocation";
  manifest.categories = [
    "System"
    "Network"
  ];
  manifest.readme = builtins.readFile ./README.md;

  # Peer options and configuration
  roles.peer = {
    description = "A peer that connects to one or more controllers.";
    interface =
      { lib, ... }:
      {
        imports = [ sharedInterface ];

        options.controller = lib.mkOption {
          type = lib.types.str;
          example = "controller1";
          description = ''
            Machinename of the controller to attach to
          '';
        };
      };

    perInstance =
      {
        instanceName,
        settings,
        roles,
        machine,
        ...
      }:
      {
        # Set default domain to instanceName

        # Peers connect to all controllers
        nixosModule =
          {
            config,
            pkgs,
            lib,
            ...
          }:
          {
            imports = [
              (extraHostsModule {
                inherit
                  instanceName
                  settings
                  roles
                  config
                  lib
                  ;
              })
            ];
            # Network allocation generator for this peer - generates host suffix
            clan.core.vars.generators."wireguard-network-${instanceName}" = {
              files.suffix.secret = false;

              runtimeInputs = with pkgs; [
                python3
              ];

              # Invalidate on hostname changes
              validation.hostname = machine.name;

              script = ''
                ${pkgs.python3}/bin/python3 ${./ipv6_allocator.py} "$out" "${instanceName}" peer "${machine.name}"
              '';
            };

            # Single wireguard interface with multiple IPs
            networking.wireguard.interfaces."${instanceName}" = {
              ips =
                # Get this peer's suffix
                let
                  peerSuffix =
                    config.clan.core.vars.generators."wireguard-network-${instanceName}".files.suffix.value;
                in
                # Create an IP in each controller's subnet
                lib.mapAttrsToList (
                  ctrlName: _:
                  let
                    controllerPrefix = builtins.readFile (
                      config.clan.core.settings.directory
                      + "/vars/per-machine/${ctrlName}/wireguard-network-${instanceName}/prefix/value"
                    );
                    peerIP = controllerPrefix + ":" + peerSuffix;
                  in
                  "${peerIP}/56"
                ) roles.controller.machines;

              privateKeyFile =
                config.clan.core.vars.generators."wireguard-keys-${instanceName}".files."privatekey".path;

              # Connect to all controllers
              peers = lib.mapAttrsToList (name: value: {
                publicKey = (
                  builtins.readFile (
                    config.clan.core.settings.directory
                    + "/vars/per-machine/${name}/wireguard-keys-${instanceName}/publickey/value"
                  )
                );

                # Allow each controller's /56 subnet
                allowedIPs = [
                  "${
                    builtins.readFile (
                      config.clan.core.settings.directory
                      + "/vars/per-machine/${name}/wireguard-network-${instanceName}/prefix/value"
                    )
                  }::/56"
                ];

                endpoint = "${value.settings.endpoint}:${toString value.settings.port}";

                persistentKeepalive = 25;
              }) roles.controller.machines;
            };
          };
      };
  };

  # Controller options and configuration
  roles.controller = {
    description = "A controller that routes peer traffic. Must be publicly reachable.";
    interface =
      { lib, ... }:
      {
        imports = [ sharedInterface ];

        options = {
          endpoint = lib.mkOption {
            type = lib.types.str;
            example = "vpn.clan.lol";
            description = ''
              Endpoint where the controller can be reached
            '';
          };
          externalPeers = lib.mkOption {
            type = lib.types.listOf lib.types.str;
            default = [ ];
            example = [
              "moms-phone"
              "daddies-laptop"
            ];
            description = ''
              List of external peer names that are not part of the clan.

              For ever entry here, a key pair for an external device will be generated.
              Each external peer must be configured on exactly ONE controller.
              This key pair can then then be displayed via `clan vars get` and inserted into an external device, like a phone or laptop.

              The names in this list must not collide with machine names in the clan.
              The machines which are part of the clan will be able to resolve the external peers via their host names, but not vice versa.
              External peers can still reach machines from within the clan via their IPv6 addresses.
            '';
          };
        };
      };
    perInstance =
      {
        settings,
        instanceName,
        roles,
        machine,
        ...
      }:
      {

        # Controllers connect to all peers and other controllers
        nixosModule =
          {
            config,
            pkgs,
            lib,
            ...
          }:
          let
            allOtherControllers = lib.filterAttrs (name: _v: name != machine.name) roles.controller.machines;
            allPeers = roles.peer.machines;
          in
          {
            imports = [
              (extraHostsModule {
                inherit
                  instanceName
                  settings
                  roles
                  config
                  lib
                  ;
              })
            ];
            # Network prefix allocation generator for this controller
            clan.core.vars.generators = {
              "wireguard-network-${instanceName}" = {
                files.prefix.secret = false;

                runtimeInputs = with pkgs; [
                  python3
                ];

                # Invalidate on network or hostname changes
                validation.hostname = machine.name;

                script = ''
                  ${pkgs.python3}/bin/python3 ${./ipv6_allocator.py} "$out" "${instanceName}" controller "${machine.name}"
                '';
              };
            }
            # For external peers, generate: suffix, public key, private key
            // lib.genAttrs' settings.externalPeers (peer: {
              name = "wireguard-network-${instanceName}-external-peer-${peer}";
              value = {
                files.suffix.secret = false;
                files.publickey.secret = false;
                files.privatekey.secret = true;
                files.privatekey.deploy = false;

                # The external peers keys are not deployed and are globally unique.
                # Even if an external peer is connected to more than one controller,
                # its private keys will remain the same.
                share = true;

                runtimeInputs = with pkgs; [
                  python3
                  wireguard-tools
                ];

                # Invalidate on hostname changes
                validation.hostname = peer;

                script = ''
                  ${pkgs.python3}/bin/python3 ${./ipv6_allocator.py} "$out" "${instanceName}" peer "${peer}"
                  wg genkey > $out/privatekey
                  wg pubkey < $out/privatekey > $out/publickey
                '';
              };
            });

            # Enable ip forwarding, so wireguard peers can reach each other
            boot.kernel.sysctl."net.ipv6.conf.all.forwarding" = 1;

            networking.firewall.allowedUDPPorts = [ settings.port ];

            # Single wireguard interface
            networking.wireguard.interfaces."${instanceName}" = {
              listenPort = settings.port;

              ips = [
                # Controller uses ::1 in its /56 subnet but with /40 prefix for proper routing
                "${config.clan.core.vars.generators."wireguard-network-${instanceName}".files.prefix.value}::1/40"
              ];

              privateKeyFile =
                config.clan.core.vars.generators."wireguard-keys-${instanceName}".files."privatekey".path;

              # Connect to all peers and other controllers
              peers =
                # Peers configuration
                (lib.mapAttrsToList (name: _value: {
                  publicKey = (
                    builtins.readFile (
                      config.clan.core.settings.directory
                      + "/vars/per-machine/${name}/wireguard-keys-${instanceName}/publickey/value"
                    )
                  );

                  # Allow the peer's /96 range in ALL controller subnets
                  allowedIPs = lib.mapAttrsToList (
                    ctrlName: _:
                    let
                      controllerPrefix = builtins.readFile (
                        config.clan.core.settings.directory
                        + "/vars/per-machine/${ctrlName}/wireguard-network-${instanceName}/prefix/value"
                      );
                      peerSuffix = builtins.readFile (
                        config.clan.core.settings.directory
                        + "/vars/per-machine/${name}/wireguard-network-${instanceName}/suffix/value"
                      );
                    in
                    "${controllerPrefix}:${peerSuffix}/96"
                  ) roles.controller.machines;

                  persistentKeepalive = 25;
                }) allPeers)
                ++
                  # External peers configuration
                  (map (peer: {
                    publicKey = (
                      builtins.readFile (
                        config.clan.core.settings.directory
                        + "/vars/shared/wireguard-network-${instanceName}-external-peer-${peer}/publickey/value"
                      )
                    );

                    # Allow the external peer's /96 range in ALL controller subnets
                    allowedIPs = lib.mapAttrsToList (
                      ctrlName: _:
                      let
                        controllerPrefix = builtins.readFile (
                          config.clan.core.settings.directory
                          + "/vars/per-machine/${ctrlName}/wireguard-network-${instanceName}/prefix/value"
                        );
                        peerSuffix = builtins.readFile (
                          config.clan.core.settings.directory
                          + "/vars/shared/wireguard-network-${instanceName}-external-peer-${peer}/suffix/value"
                        );
                      in
                      "${controllerPrefix}:${peerSuffix}/96"
                    ) roles.controller.machines;

                    # No endpoint for external peers, they initiate the connection
                    persistentKeepalive = 25;
                  }) settings.externalPeers)
                ++
                  # Other controllers configuration
                  (lib.mapAttrsToList (name: value: {
                    publicKey = (
                      builtins.readFile (
                        config.clan.core.settings.directory
                        + "/vars/per-machine/${name}/wireguard-keys-${instanceName}/publickey/value"
                      )
                    );

                    allowedIPs = [
                      "${
                        builtins.readFile (
                          config.clan.core.settings.directory
                          + "/vars/per-machine/${name}/wireguard-network-${instanceName}/prefix/value"
                        )
                      }::/56"
                    ];

                    endpoint = "${value.settings.endpoint}:${toString value.settings.port}";
                    persistentKeepalive = 25;
                  }) allOtherControllers);
            };
          };
      };
  };

  # Maps over all machines and produces one result per machine, regardless of role
  perMachine =
    { instances, machine, ... }:
    {
      nixosModule =
        { pkgs, lib, ... }:
        let
          # Check if this machine has conflicting roles across all instances
          machineRoleConflicts = lib.flatten (
            lib.mapAttrsToList (
              instanceName: instanceInfo:
              let
                isController =
                  instanceInfo.roles ? controller && instanceInfo.roles.controller.machines ? ${machine.name};
                isPeer = instanceInfo.roles ? peer && instanceInfo.roles.peer.machines ? ${machine.name};
              in
              lib.optional (isController && isPeer) {
                inherit instanceName;
                machineName = machine.name;
              }
            ) instances
          );
        in
        {
          # Add assertions for role conflicts
          assertions = lib.forEach machineRoleConflicts (conflict: {
            assertion = false;
            message = ''
              Machine '${conflict.machineName}' cannot have both 'controller' and 'peer' roles in the wireguard instance '${conflict.instanceName}'.
              A machine must be either a controller or a peer, not both.
            '';
          });

          # Generate keys for each instance where this machine participates
          clan.core.vars.generators = lib.mapAttrs' (
            name: _instanceInfo:
            lib.nameValuePair "wireguard-keys-${name}" {
              files.publickey.secret = false;
              files.privatekey = { };

              runtimeInputs = with pkgs; [
                wireguard-tools
              ];

              script = ''
                wg genkey > $out/privatekey
                wg pubkey < $out/privatekey > $out/publickey
              '';
            }
          ) instances;

        };
    };
}