{ config, lib, pkgs, ... }:

with lib;

let
  cfg = config.nix-nos.vlans;

  vlanType = types.submodule {
    options = {
      id = mkOption {
        type = types.int;
        description = "VLAN ID (1-4094)";
        example = 100;
      };

      interface = mkOption {
        type = types.str;
        description = "Parent interface";
        example = "eth0";
      };

      addresses = mkOption {
        type = types.listOf types.str;
        default = [];
        description = "IP addresses for this VLAN interface (CIDR notation)";
        example = [ "10.0.100.1/24" ];
      };

      gateway = mkOption {
        type = types.nullOr types.str;
        default = null;
        description = "Default gateway for this VLAN";
        example = "10.0.100.254";
      };

      dns = mkOption {
        type = types.listOf types.str;
        default = [];
        description = "DNS servers for this VLAN";
        example = [ "8.8.8.8" ];
      };

      mtu = mkOption {
        type = types.nullOr types.int;
        default = null;
        description = "MTU size for this VLAN interface";
        example = 1500;
      };
    };
  };

in {
  options.nix-nos.vlans = mkOption {
    type = types.attrsOf vlanType;
    default = {};
    description = "VLAN configurations using systemd-networkd";
    example = literalExpression ''
      {
        storage = {
          id = 100;
          interface = "eth0";
          addresses = [ "10.0.100.1/24" ];
        };
        mgmt = {
          id = 200;
          interface = "eth0";
          addresses = [ "10.0.200.1/24" ];
          gateway = "10.0.200.254";
        };
      }
    '';
  };

  config = mkIf (cfg != {}) {
    assertions = [
      {
        assertion = all (name: vlan:
          vlan.id >= 1 && vlan.id <= 4094
        ) (mapAttrsToList nameValuePair cfg);
        message = "nix-nos.vlans: VLAN ID must be between 1 and 4094";
      }
      {
        assertion = all (name: vlan:
          (length vlan.addresses) > 0
        ) (mapAttrsToList nameValuePair cfg);
        message = "nix-nos.vlans: Each VLAN must have at least one address";
      }
    ];

    systemd.network.enable = true;

    # Create VLAN netdevs
    systemd.network.netdevs = mapAttrs' (name: vlan:
      nameValuePair "20-${name}" {
        netdevConfig = {
          Name = name;
          Kind = "vlan";
        };
        vlanConfig.Id = vlan.id;
      }) cfg;

    # Configure VLAN networks
    systemd.network.networks = mkMerge [
      # VLAN interface networks
      (mapAttrs' (name: vlan:
        nameValuePair "20-${name}" {
          matchConfig.Name = name;
          address = vlan.addresses;
          gateway = optional (vlan.gateway != null) vlan.gateway;
          dns = vlan.dns;
          linkConfig = optionalAttrs (vlan.mtu != null) {
            MTUBytes = toString vlan.mtu;
          };
        }) cfg)

      # Parent interface VLAN attachment
      # Group VLANs by parent interface
      (let
        vlansByParent = foldl' (acc: nameVlanPair:
          let
            name = nameVlanPair.name;
            vlan = nameVlanPair.value;
            parent = vlan.interface;
          in
          acc // {
            ${parent} = (acc.${parent} or []) ++ [ name ];
          }
        ) {} (mapAttrsToList nameValuePair cfg);
      in
        mapAttrs' (parent: vlanNames:
          nameValuePair "21-${parent}-vlans" {
            matchConfig.Name = parent;
            vlan = vlanNames;
          }) vlansByParent)
    ];
  };
}