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

with lib;

let
  cfg = config.plasmacloud.cluster;

  # 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 = "Services to run: chainfire, flaredb, iam, etc.";
        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";
      };
    };
  };

  # Generate cluster-config.json for the current node
  generateClusterConfig = cluster:
    let
      hostname = config.networking.hostName;
      node = cluster.nodes.${hostname} or (throw "Node ${hostname} not found in cluster configuration");

      # Determine bootstrap node (first node in initialPeers list)
      bootstrapNodeName = head cluster.bootstrap.initialPeers;
      isBootstrap = hostname == bootstrapNodeName;

      # Get bootstrap node config
      bootstrapNode = cluster.nodes.${bootstrapNodeName};

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

      # Control plane nodes for Raft peers
      controlPlaneNodes = filter (n: cluster.nodes.${n}.role == "control-plane") (attrNames cluster.nodes);

      # 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;
      bgp_asn = cluster.bgp.asn;
    };

in {
  options.plasmacloud.cluster = {
    enable = mkEnableOption "PlasmaCloud cluster configuration";

    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" "iam" ];
          };
        }
      '';
    };

    bootstrap = {
      initialPeers = mkOption {
        type = types.listOf types.str;
        description = "Initial Raft peers for bootstrap (ordered list, first is bootstrap node)";
        example = [ "node01" "node02" "node03" ];
      };
    };

    bgp = {
      asn = mkOption {
        type = types.int;
        description = "BGP AS number for the cluster";
        example = 64512;
      };
    };
  };

  config = mkIf cfg.enable {
    # Assertions
    assertions = [
      {
        assertion = (length (attrNames cfg.nodes)) > 0;
        message = "plasmacloud.cluster.nodes must contain at least one node";
      }
      {
        assertion = (length cfg.bootstrap.initialPeers) > 0;
        message = "plasmacloud.cluster.bootstrap.initialPeers must contain at least one node";
      }
      {
        assertion = all (peer: cfg.nodes ? ${peer}) cfg.bootstrap.initialPeers;
        message = "All nodes in bootstrap.initialPeers must exist in cluster.nodes";
      }
      {
        assertion = cfg.bgp.asn > 0 && cfg.bgp.asn < 4294967296;
        message = "BGP ASN must be between 1 and 4294967295";
      }
    ];

    # Generate cluster-config.json for first-boot-automation
    environment.etc."nixos/secrets/cluster-config.json" = mkIf (cfg.nodes ? ${config.networking.hostName}) {
      text = builtins.toJSON (generateClusterConfig cfg);
      mode = "0600";
    };
  };
}