# SPDX-FileCopyrightText: 2024 Luke Granger-Brown <depot@lukegb.com>
#
# SPDX-License-Identifier: Apache-2.0

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

let
  cfg = config.my.services.open5gs;

  settingsFormat = pkgs.formats.yaml {};
in
{
  options.my.services.open5gs = let
    package = lib.mkOption {
      type = lib.types.package;
      default = depot.nix.pkgs.open5gs;
    };
    settings = lib.mkOption {
      type = lib.types.submodule {
        freeformType = settingsFormat.type;
      };
      default = {};
    };
  in {
    enable = lib.mkEnableOption "Open5GS";
    inherit package;
    globalSettings = settings;
    enableSAComponents = (lib.mkEnableOption "default Open5GS components for 5G SA core") // {
      default = true;
    };
    enableNSAComponents = (lib.mkEnableOption "default Open5GS components for 4G/5G NSA core") // {
      default = true;
    };
    webui = {
      enable = lib.mkEnableOption "Open5GS WebUI";
      package = lib.mkOption {
        type = lib.types.package;
        default = depot.nix.pkgs.open5gs.webui;
      };

      secretKeyFile = lib.mkOption {
        type = lib.types.path;
      };
      databaseURI = lib.mkOption {
        type = lib.types.str;
        default = "mongodb://localhost/open5gs";
      };
      listenHostname = lib.mkOption {
        type = lib.types.str;
        default = "localhost";
      };
      listenPort = lib.mkOption {
        type = lib.types.port;
        default = 9999;
      };
    };

    nrf = {
      enable = lib.mkEnableOption "Open5GS NF Repository Function";
      inherit package settings;
    };
    scp = {
      enable = lib.mkEnableOption "Open5GS Service Communication Proxy";
      inherit package settings;
    };
    sepp = {
      enable = lib.mkEnableOption "Open5GS Security Edge Protection Proxy";
      inherit package settings;
    };
    amf = {
      enable = lib.mkEnableOption "Open5GS Access and Mobility Management Function";
      inherit package settings;
    };
    smf = {
      enable = lib.mkEnableOption "Open5GS Session Management Function";
      inherit package settings;
    };
    upf = {
      enable = lib.mkEnableOption "Open5GS User Plane Function";
      inherit package settings;
    };
    ausf = {
      enable = lib.mkEnableOption "Open5GS Authentication Server Function";
      inherit package settings;
    };
    udm = {
      enable = lib.mkEnableOption "Open5GS Unified Data Management";
      inherit package settings;
    };
    udr = {
      enable = lib.mkEnableOption "Open5GS Unified Data Repository";
      inherit package settings;
    };
    pcf = {
      enable = lib.mkEnableOption "Policy and Charging Function";
      inherit package settings;
    };
    nssf = {
      enable = lib.mkEnableOption "Network Slice Selection Function";
      inherit package settings;
    };
    bsf = {
      enable = lib.mkEnableOption "Binding Support Function";
      inherit package settings;
    };

    mme = {
      enable = lib.mkEnableOption "Mobility Management Entity";
      inherit package settings;
    };
    hss = {
      enable = lib.mkEnableOption "Home Subscriber Server";
      inherit package settings;
    };
    pcrf = {
      enable = lib.mkEnableOption "Policy and Charging Rules Function";
      inherit package settings;
    };
    sgwc = {
      enable = lib.mkEnableOption "Serving Gateway Control Plane";
      inherit package settings;
    };
    sgwu = {
      enable = lib.mkEnableOption "Serving Gateway User Plane";
      inherit package settings;
    };
  };

  config = let
    mkComponent = componentCfg: componentName: defaultSettings: lib.mkIf componentCfg.enable {
      my.services.open5gs."${componentName}".settings = lib.mkMerge [cfg.globalSettings {
        logger.file.path = lib.mkDefault "/var/log/open5gs/${componentName}.log";
      } defaultSettings];
      environment.etc."open5gs/${componentName}.yaml".source = settingsFormat.generate "open5gs-${componentName}-config.yaml" componentCfg.settings;
      systemd.services."open5gs-${componentName}" = {
        wantedBy = [ "open5gs.target" ];
        restartTriggers = [ config.environment.etc."open5gs/${componentName}.yaml".source or null ];
        serviceConfig = {
          LogsDirectory = "open5gs";
          ExecStart = "${componentCfg.package}/bin/open5gs-${componentName}d -c /etc/open5gs/${componentName}.yaml";
        };
      };
    };
  in lib.mkIf cfg.enable (lib.mkMerge [
    {
      services.mongodb.enable = lib.mkDefault true;
      my.services.open5gs = {
        nrf.enable = lib.mkDefault cfg.enableSAComponents;
        nrf.package = lib.mkDefault cfg.package;
        scp.enable = lib.mkDefault cfg.enableSAComponents;
        scp.package = lib.mkDefault cfg.package;
        sepp.enable = lib.mkDefault cfg.enableSAComponents;
        sepp.package = lib.mkDefault cfg.package;
        amf.enable = lib.mkDefault cfg.enableSAComponents;
        amf.package = lib.mkDefault cfg.package;
        smf.enable = lib.mkDefault (cfg.enableSAComponents || cfg.enableNSAComponents);
        smf.package = lib.mkDefault cfg.package;
        upf.enable = lib.mkDefault (cfg.enableSAComponents || cfg.enableNSAComponents);
        upf.package = lib.mkDefault cfg.package;
        ausf.enable = lib.mkDefault cfg.enableSAComponents;
        ausf.package = lib.mkDefault cfg.package;
        udm.enable = lib.mkDefault cfg.enableSAComponents;
        udm.package = lib.mkDefault cfg.package;
        udr.enable = lib.mkDefault cfg.enableSAComponents;
        udr.package = lib.mkDefault cfg.package;
        pcf.enable = lib.mkDefault cfg.enableSAComponents;
        pcf.package = lib.mkDefault cfg.package;
        nssf.enable = lib.mkDefault cfg.enableSAComponents;
        nssf.package = lib.mkDefault cfg.package;
        bsf.enable = lib.mkDefault cfg.enableSAComponents;
        bsf.package = lib.mkDefault cfg.package;
        mme.enable = lib.mkDefault cfg.enableNSAComponents;
        mme.package = lib.mkDefault cfg.package;
        hss.enable = lib.mkDefault cfg.enableNSAComponents;
        hss.package = lib.mkDefault cfg.package;
        pcrf.enable = lib.mkDefault cfg.enableNSAComponents;
        pcrf.package = lib.mkDefault cfg.package;
        sgwc.enable = lib.mkDefault cfg.enableNSAComponents;
        sgwc.package = lib.mkDefault cfg.package;
        sgwu.enable = lib.mkDefault cfg.enableNSAComponents;
        sgwu.package = lib.mkDefault cfg.package;
        webui.enable = lib.mkDefault (cfg.enableSAComponents || cfg.enableNSAComponents);

        globalSettings = {
          db_uri = lib.mkDefault "mongodb://localhost/open5gs";
          logger.level = lib.mkDefault "info";
          global = {
            max.ue = lib.mkDefault 1024;
            parameter = {
              no_nrf = lib.mkDefault (!cfg.enableSAComponents);
              no_scp = lib.mkDefault (!cfg.enableSAComponents);
              no_sepp = lib.mkDefault (!cfg.enableSAComponents);
              no_amf = lib.mkDefault (!cfg.enableSAComponents);
              no_smf = lib.mkDefault (!cfg.enableSAComponents && !cfg.enableNSAComponents);
              no_upf = lib.mkDefault (!cfg.enableSAComponents && !cfg.enableNSAComponents);
              no_ausf = lib.mkDefault (!cfg.enableSAComponents);
              no_udm = lib.mkDefault (!cfg.enableSAComponents);
              no_udr = lib.mkDefault (!cfg.enableSAComponents);
              no_pcf = lib.mkDefault (!cfg.enableSAComponents);
              no_nssf = lib.mkDefault (!cfg.enableSAComponents);
              no_bsf = lib.mkDefault (!cfg.enableSAComponents);

              no_mme = lib.mkDefault (!cfg.enableNSAComponents);
              no_sgwc = lib.mkDefault (!cfg.enableNSAComponents);
              no_sgwu = lib.mkDefault (!cfg.enableNSAComponents);
              no_pcrf = lib.mkDefault (!cfg.enableNSAComponents);
              no_hss = lib.mkDefault (!cfg.enableNSAComponents);
            };
          };
        };
      };

      users.groups.open5gs = {};

      systemd.targets.open5gs = {
        description = "Open5GS";
        wantedBy = [ "multi-user.target" ];
      };
    }
    (lib.mkIf cfg.webui.enable {
      systemd.services.open5gs-webui = {
        wantedBy = [ "open5gs.target" ];
        script = ''
          export SECRET_KEY="$(cat "''${CREDENTIALS_DIRECTORY}/secret_key")"
          exec ${cfg.webui.package}/bin/open5gs-webui
        '';
        serviceConfig = {
          LoadCredential = "secret_key:${cfg.webui.secretKeyFile}";
          DynamicUser = true;
          User = "open5gs-webui";
          Group = "open5gs";
        };
        environment = {
          SECRET_KEY = cfg.webui.secretKeyFile;
          DB_URI = cfg.webui.databaseURI;
          HOSTNAME = cfg.webui.listenHostname;
          PORT = toString cfg.webui.listenPort;
        };
      };
    })
    (mkComponent cfg.nrf "nrf" {
      nrf = lib.mkDefault {
        serving = [
          {
            plmn_id = {
              mcc = 999;
              mnc = 42;
            };
          }
        ];
        sbi.server = [
          {
            address = "127.0.0.10";
            port = 7777;
          }
        ];
      };
    })
    (mkComponent cfg.scp "scp" {
      scp = lib.mkDefault {
        sbi = {
          server = [
            {
              address = "127.0.0.200";
              port = 7777;
            }
          ];
          client = {
            nrf = [
              {
                uri = "http://127.0.0.10:7777";
              }
            ];
          };
        };
      };
    })
    (mkComponent cfg.sepp "sepp" {
      sepp = lib.mkDefault {
        default.tls = {
          server = {
            private_key = "${cfg.sepp.package}/etc/open5gs/tls/sepp1.key";
            cert = "${cfg.sepp.package}/etc/open5gs/tls/sepp1.crt";
          };
          client = {
            cacert = "${cfg.sepp.package}/etc/open5gs/tls/ca.crt";
          };
        };
        sbi = {
          server = [
            {
              address = "127.0.1.250";
              port = 7777;
            }
          ];
          client = {
            scp = [
              {
                uri = "http://127.0.0.200:7777";
              }
            ];
          };
        };
        n32 = {
          server = [
            {
              sender = "sepp1.localdomain";
              scheme = "https";
              address = "127.0.1.251";
              port = 7777;
              n32f = {
                scheme = "https";
                address = "127.0.1.252";
                port = 7777;
              };
            }
          ];
        };
      };
    })
    (mkComponent cfg.amf "amf" {
      amf = lib.mkDefault {
        sbi = {
          server = [
            {
              address = "127.0.0.5";
              port = 7777;
            }
          ];
          client = {
            scp = [
              {
                uri = "http://127.0.0.200:7777";
              }
            ];
          };
        };
        ngap.server = [
          {
            address = "127.0.0.5";
          }
        ];
        metrics.server = [
          {
            address = "127.0.0.5";
            port = 9090;
          }
        ];
        guami = [
          {
            plmn_id = {
              mcc = 999;
              mnc = 42;
            };
            amf_id = {
              region = 2;
              set = 1;
            };
          }
        ];
        tai = [
          {
            plmn_id = {
              mcc = 999;
              mnc = 42;
            };
            tac = 1;
          }
        ];
        plmn_support = [
          {
            plmn_id = {
              mcc = 999;
              mnc = 42;
            };
            s_nssai = [
              {
                sst = 1;
              }
            ];
          }
        ];
        security = {
          integrity_order = [
            "NIA2"
            "NIA1"
            "NIA0"
          ];
          ciphering_order = [
            "NEA0"
            "NEA1"
            "NEA2"
          ];
        };
        network_name = {
          full = "Open5GS";
          short = "Next";
        };
        amf_name = "open5gs-amf0";
        time.t3512.value = 540;
      };
    })
    (mkComponent cfg.smf "smf" {
      smf = lib.mkDefault {
        sbi = {
          server = [
            {
              address = "127.0.0.4";
              port = 7777;
            }
          ];
          client = {
            scp = [
              {
                uri = "http://127.0.0.200:7777";
              }
            ];
          };
        };
        pfcp = {
          server = [
            {
              address = "127.0.0.4";
            }
          ];
          client = {
            upf = [
              {
                address = "127.0.0.7";
              }
            ];
          };
        };
        gtpc.server = [
          {
            address = "127.0.0.4";
          }
        ];
        gtpu.server = [
          {
            address = "127.0.0.4";
          }
        ];
        metrics.server = [
          {
            address = "127.0.0.4";
            port = 9090;
          }
        ];
        session = [
          {
            subnet = "10.45.0.0/16";
            gateway = "10.45.0.1";
          }
          {
            subnet = "2001:db8:cafe::/48";
            gateway = "2001:db8:cafe::1";
          }
        ];
        dns = [
          "8.8.8.8"
          "8.8.4.4"
          "2001:4860:4860::8888"
          "2001:4860:4860::8844"
        ];
        mtu = 1400;
        freeDiameter = "${cfg.smf.package}/etc/freeDiameter/smf.conf";
      };
    })
    (mkComponent cfg.upf "upf" {
      upf = lib.mkDefault {
        pfcp = {
          server = [
            {
              address = "127.0.0.7";
            }
          ];
          client = null;
        };
        gtpu.server = [
          {
            address = "127.0.0.7";
          }
        ];
        session = [
          {
            subnet = "10.45.0.0/16";
            gateway = "10.45.0.1";
          }
          {
            subnet = "2001:db8:cafe::/48";
            gateway = "2001:db8:cafe::1";
          }
        ];
        metrics.server = [
          {
            address = "127.0.0.7";
            port = 9090;
          }
        ];
      };
    })
    (mkComponent cfg.ausf "ausf" {
      ausf = lib.mkDefault {
        sbi = {
          server = [
            {
              address = "127.0.0.11";
              port = 7777;
            }
          ];
          client = {
            scp = [
              {
                uri = "http://127.0.0.200:7777";
              }
            ];
          };
        };
      };
    })
    (mkComponent cfg.udm "udm" {
      udm = lib.mkDefault {
        hnet = [
          {
            id = 1;
            scheme = 1;
            key = "${cfg.udm.package}/etc/open5gs/hnet/curve25519-1.key";
          }
          {
            id = 2;
            scheme = 2;
            key = "${cfg.udm.package}/etc/open5gs/hnet/secp256r1-2.key";
          }
          {
            id = 3;
            scheme = 1;
            key = "${cfg.udm.package}/etc/open5gs/hnet/curve25519-3.key";
          }
          {
            id = 4;
            scheme = 2;
            key = "${cfg.udm.package}/etc/open5gs/hnet/secp256r1-4.key";
          }
          {
            id = 5;
            scheme = 1;
            key = "${cfg.udm.package}/etc/open5gs/hnet/curve25519-5.key";
          }
          {
            id = 6;
            scheme = 2;
            key = "${cfg.udm.package}/etc/open5gs/hnet/secp256r1-6.key";
          }
        ];
        sbi = {
          server = [
            {
              address = "127.0.0.12";
              port = 7777;
            }
          ];
          client = {
            scp = [
              {
                uri = "http://127.0.0.200:7777";
              }
            ];
          };
        };
      };
    })
    (mkComponent cfg.udr "udr" {
      udr = lib.mkDefault {
        sbi = {
          server = [
            {
              address = "127.0.0.20";
              port = 7777;
            }
          ];
          client = {
            scp = [
              {
                uri = "http://127.0.0.200:7777";
              }
            ];
          };
        };
      };
    })
    (mkComponent cfg.pcf "pcf" {
      pcf = lib.mkDefault {
        sbi = {
          server = [
            {
              address = "127.0.0.13";
              port = 7777;
            }
          ];
          client = {
            scp = [
              {
                uri = "http://127.0.0.200:7777";
              }
            ];
          };
        };
        metrics.server = [
          {
            address = "127.0.0.13";
            port = 9090;
          }
        ];
      };
    })
    (mkComponent cfg.nssf "nssf" {
      nssf = lib.mkDefault {
        sbi = {
          server = [
            {
              address = "127.0.0.14";
              port = 7777;
            }
          ];
          client = {
            scp = [
              {
                uri = "http://127.0.0.200:7777";
              }
            ];
            nsi = [
              {
                uri = "http://127.0.0.10:7777";
                s_nssai = {
                  sst = 1;
                };
              }
            ];
          };
        };
      };
    })
    (mkComponent cfg.bsf "bsf" {
      bsf = lib.mkDefault {
        sbi = {
          server = [
            {
              address = "127.0.0.15";
              port = 7777;
            }
          ];
          client = {
            scp = [
              {
                uri = "http://127.0.0.200:7777";
              }
            ];
          };
        };
      };
    })

    (mkComponent cfg.mme "mme" {
      mme = lib.mkDefault {
        freeDiameter = "${cfg.mme.package}/etc/freeDiameter/mme.conf";
        s1ap.server = [
          {
            address = "127.0.0.2";
          }
        ];
        gtpc = {
          server = [
            {
              address = "127.0.0.2";
            }
          ];
          client = {
            sgwc = [
              {
                address = "127.0.0.3";
              }
            ];
            smf = [
              {
                address = "127.0.0.4";
              }
            ];
          };
        };
        metrics.server = [
          {
            address = "127.0.0.2";
            port = 9090;
          }
        ];
        gummei = [
          {
            plmn_id = {
              mcc = 999;
              mnc = 42;
            };
            mme_gid = 2;
            mme_code = 1;
          }
        ];
        tai = [
          {
            plmn_id = {
              mcc = 999;
              mnc = 42;
            };
            tac = 1;
          }
        ];
        security = {
          integrity_order = [
            "EIA2"
            "EIA1"
            "EIA0"
          ];
          ciphering_order = [
            "EEA0"
            "EEA1"
            "EEA2"
          ];
        };
        network_name = {
          full = "Open5GS";
          short = "Next";
        };
        mme_name = "open5gs-mme0";
        time = null;
      };
    })
    (mkComponent cfg.hss "hss" {
      hss = lib.mkDefault {
        freeDiameter = "${cfg.hss.package}/etc/freeDiameter/hss.conf";
      };
    })
    (mkComponent cfg.pcrf "pcrf" {
      pcrf = lib.mkDefault {
        freeDiameter = "${cfg.pcrf.package}/etc/freeDiameter/pcrf.conf";
      };
    })
    (mkComponent cfg.sgwc "sgwc" {
      sgwc = lib.mkDefault {
        gtpc.server = [
          {
            address = "127.0.0.3";
          }
        ];
        pfcp = {
          server = [
            {
              address = "127.0.0.3";
            }
          ];
          client = {
            sgwu = [
              {
                address = "127.0.0.6";
              }
            ];
          };
        };
      };
    })
    (mkComponent cfg.sgwu "sgwu" {
      sgwu = lib.mkDefault {
        pfcp = {
          server = [
            {
              address = "127.0.0.6";
            }
          ];
          client = null;
        };
        gtpu.server = [
          {
            address = "127.0.0.6";
          }
        ];
      };
    })
  ]);
}