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

with lib;

let
  cfg = config.nix-nos;

  # Node definition type
  nodeType = types.submodule {
    options = {
      role = mkOption {
        type = types.enum [ "control-plane" "worker" ];
        default = "worker";
        description = "Node role in the cluster";
      };

      ip = mkOption {
        type = types.str;
        description = "IP address of the node";
      };

      services = mkOption {
        type = types.listOf types.str;
        default = [];
        description = "List of services to run on this node";
        example = [ "chainfire" "flaredb" "iam" ];
      };

      raftPort = mkOption {
        type = types.port;
        default = 2380;
        description = "Raft port for consensus protocols";
      };

      apiPort = mkOption {
        type = types.port;
        default = 2379;
        description = "API port for cluster services";
      };

      metadata = mkOption {
        type = types.attrsOf types.anything;
        default = {};
        description = "Additional metadata for the node";
      };
    };
  };

  # Cluster definition type
  clusterType = types.submodule {
    options = {
      name = mkOption {
        type = types.str;
        default = "plasmacloud-cluster";
        description = "Cluster name";
      };

      nodes = mkOption {
        type = types.attrsOf nodeType;
        default = {};
        description = "Map of node names to their configurations";
        example = literalExpression ''
          {
            "node01" = {
              role = "control-plane";
              ip = "10.0.1.10";
              services = [ "chainfire" "flaredb" ];
            };
          }
        '';
      };

      bootstrapNode = mkOption {
        type = types.nullOr types.str;
        default = null;
        description = "Name of the bootstrap node (first control-plane node if null)";
      };
    };
  };

in {
  options.nix-nos = {
    enable = mkEnableOption "Nix-NOS declarative cluster management";

    clusters = mkOption {
      type = types.attrsOf clusterType;
      default = {};
      description = "Map of cluster names to their configurations";
    };

    # Helper function to generate cluster-config.json for a specific node
    generateClusterConfig = mkOption {
      type = types.functionTo types.attrs;
      default = { hostname, clusterName ? "plasmacloud" }:
        let
          cluster = cfg.clusters.${clusterName} or (throw "Cluster ${clusterName} not found");
          node = cluster.nodes.${hostname} or (throw "Node ${hostname} not found in cluster ${clusterName}");

          # Determine bootstrap node
          controlPlaneNodes = filter (n: cluster.nodes.${n}.role == "control-plane") (attrNames cluster.nodes);
          bootstrapNodeName =
            if cluster.bootstrapNode != null
            then cluster.bootstrapNode
            else head controlPlaneNodes;

          isBootstrap = hostname == bootstrapNodeName;

          # Leader URL (bootstrap node's API endpoint)
          bootstrapNode = cluster.nodes.${bootstrapNodeName};
          leaderUrl = "https://${bootstrapNode.ip}:${toString bootstrapNode.apiPort}";

          # Initial peers for Raft cluster
          initialPeers = map (nodeName: {
            id = nodeName;
            addr = "${cluster.nodes.${nodeName}.ip}:${toString cluster.nodes.${nodeName}.raftPort}";
          }) controlPlaneNodes;

          # FlareDB peers (all control-plane nodes)
          flaredbPeers = map (nodeName:
            "${cluster.nodes.${nodeName}.ip}:${toString (cluster.nodes.${nodeName}.apiPort + 100)}"
          ) controlPlaneNodes;

        in {
          node_id = hostname;
          node_role = node.role;
          bootstrap = isBootstrap;
          cluster_name = cluster.name;
          leader_url = leaderUrl;
          raft_addr = "${node.ip}:${toString node.raftPort}";
          initial_peers = initialPeers;
          flaredb_peers = flaredbPeers;
          services = node.services;
          metadata = node.metadata;
        };
      description = "Function to generate cluster-config.json for a specific hostname";
    };
  };

  config = mkIf cfg.enable {
    # Ensure at least one cluster is defined
    assertions = [
      {
        assertion = (builtins.length (attrNames cfg.clusters)) > 0;
        message = "nix-nos.clusters must contain at least one cluster definition";
      }
    ];
  };
}