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

with lib;

let
  cfg = config.roles.k3s-node;
in
{
  options.roles.k3s-node = {
    enable = mkEnableOption "Enable k3s node";

    role = mkOption {
      type = types.enum [ "server" "agent" ];
      default = "agent";
      description = "k3s role: server (control plane) or agent (worker)";
    };

    serverAddr = mkOption {
      type = types.str;
      default = "https://10.0.0.222:6443";
      description = "URL of k3s server to join (required for agents, used for HA servers)";
    };

    tokenFile = mkOption {
      type = types.path;
      default = "/etc/k3s/token";
      description = "Path to file containing the cluster join token";
    };

    clusterInit = mkOption {
      type = types.bool;
      default = false;
      description = "Initialize a new cluster (first server only)";
    };

    extraFlags = mkOption {
      type = types.listOf types.str;
      default = [];
      description = "Additional flags to pass to k3s";
    };

    gracefulNodeShutdown = mkOption {
      type = types.bool;
      default = true;
      description = "Enable graceful node shutdown";
    };

    openFirewall = mkOption {
      type = types.bool;
      default = true;
      description = "Open firewall ports for k3s";
    };
  };

  config = mkIf cfg.enable {
    # k3s service configuration
    services.k3s = {
      enable = true;
      role = cfg.role;
      tokenFile = cfg.tokenFile;
      extraFlags = cfg.extraFlags;
      gracefulNodeShutdown.enable = cfg.gracefulNodeShutdown;
      serverAddr = if (cfg.role == "agent" || !cfg.clusterInit) then cfg.serverAddr else "";
      clusterInit = cfg.role == "server" && cfg.clusterInit;
    };

    # Firewall rules for k3s
    networking.firewall = mkIf cfg.openFirewall {
      allowedTCPPorts = [
        6443  # k3s API server
        10250 # kubelet metrics
      ] ++ optionals (cfg.role == "server") [
        2379  # etcd clients (HA)
        2380  # etcd peers (HA)
      ];
      allowedUDPPorts = [
        8472  # flannel VXLAN
      ];
    };
  };
}