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

with lib;

let
  cfg = config.plasmacloud.cluster;
  clusterConfigLib = import ../../nix-nos/lib/cluster-config-lib.nix { inherit lib; };
  nodeType = clusterConfigLib.mkNodeType types;
  nodeClassType = clusterConfigLib.mkNodeClassType types;
  nodePoolType = clusterConfigLib.mkNodePoolType types;
  enrollmentRuleType = clusterConfigLib.mkEnrollmentRuleType types;
  hostDeploymentType = clusterConfigLib.mkHostDeploymentType types;
  serviceType = clusterConfigLib.mkServiceType types;
  mtlsPolicyType = clusterConfigLib.mkMtlsPolicyType types;
  jsonFormat = pkgs.formats.json { };

  # Generate cluster-config.json for the current node
  generateClusterConfig = cluster:
    clusterConfigLib.mkClusterConfig {
      inherit cluster;
      hostname = config.networking.hostName;
      bootstrapNodeName = head cluster.bootstrap.initialPeers;
    };

  generatedNodeClusterConfig =
    if cfg.nodes ? "${config.networking.hostName}" then
      generateClusterConfig cfg
    else
      null;

  generatedNixNOSTopologyCluster = clusterConfigLib.mkNixNOSTopologyCluster cfg;
  generatedDeployerClusterState = clusterConfigLib.mkDeployerClusterState cfg;

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;
      };
    };

    deployer = {
      clusterId = mkOption {
        type = types.nullOr types.str;
        default = null;
        description = "Cluster ID exported into deployer cluster state";
      };

      environment = mkOption {
        type = types.nullOr types.str;
        default = null;
        description = "Optional environment label exported into deployer cluster state";
      };

      nodeClasses = mkOption {
        type = types.attrsOf nodeClassType;
        default = { };
        description = "Reusable deployer node classes derived from Nix";
      };

      pools = mkOption {
        type = types.attrsOf nodePoolType;
        default = { };
        description = "Reusable deployer node pools derived from Nix";
      };

      enrollmentRules = mkOption {
        type = types.attrsOf enrollmentRuleType;
        default = { };
        description = "Deployer auto-enrollment rules derived from Nix";
      };

      hostDeployments = mkOption {
        type = types.attrsOf hostDeploymentType;
        default = { };
        description = "Declarative host rollout objects derived from Nix";
      };

      services = mkOption {
        type = types.attrsOf serviceType;
        default = { };
        description = "Scheduler-managed service definitions derived from Nix";
      };

      mtlsPolicies = mkOption {
        type = types.attrsOf mtlsPolicyType;
        default = { };
        description = "Declarative mTLS policies derived from Nix";
      };
    };

    generated = {
      nodeClusterConfig = mkOption {
        type = types.nullOr (types.attrsOf types.anything);
        internal = true;
        default = null;
        description = "Resolved per-node cluster-config.json content";
      };

      deployerClusterState = mkOption {
        type = types.attrsOf types.anything;
        internal = true;
        default = { };
        description = "Resolved deployer cluster state exported from Nix";
      };
    };
  };

  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";
      }
      {
        assertion = all (nodeName:
          let
            node = cfg.nodes.${nodeName};
          in
            node.pool == null || cfg.deployer.pools ? "${node.pool}"
        ) (attrNames cfg.nodes);
        message = "All node pools referenced in plasmacloud.cluster.nodes must exist in plasmacloud.cluster.deployer.pools";
      }
      {
        assertion = all (nodeName:
          let
            node = cfg.nodes.${nodeName};
          in
            node.nodeClass == null || cfg.deployer.nodeClasses ? "${node.nodeClass}"
        ) (attrNames cfg.nodes);
        message = "All node classes referenced in plasmacloud.cluster.nodes must exist in plasmacloud.cluster.deployer.nodeClasses";
      }
      {
        assertion = all (poolName:
          let
            pool = cfg.deployer.pools.${poolName};
          in
            pool.nodeClass == null || cfg.deployer.nodeClasses ? "${pool.nodeClass}"
        ) (attrNames cfg.deployer.pools);
        message = "All deployer pools must reference existing deployer node classes";
      }
      {
        assertion = all (ruleName:
          let
            rule = cfg.deployer.enrollmentRules.${ruleName};
          in
            (rule.pool == null || cfg.deployer.pools ? "${rule.pool}")
            && (rule.nodeClass == null || cfg.deployer.nodeClasses ? "${rule.nodeClass}")
        ) (attrNames cfg.deployer.enrollmentRules);
        message = "All deployer enrollment rules must reference existing pools and node classes";
      }
      {
        assertion = all (deploymentName:
          let
            deployment = cfg.deployer.hostDeployments.${deploymentName};
          in
            all (pool: cfg.deployer.pools ? "${pool}") deployment.selector.pools
            && all (nodeClass: cfg.deployer.nodeClasses ? "${nodeClass}") deployment.selector.nodeClasses
        ) (attrNames cfg.deployer.hostDeployments);
        message = "All deployer host deployments must reference existing pools and node classes";
      }
    ];

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

    nix-nos.enable = mkDefault true;
    nix-nos.clusters = {
      "${cfg.name}" = mkDefault generatedNixNOSTopologyCluster;
    };

    plasmacloud.cluster.generated.nodeClusterConfig = generatedNodeClusterConfig;
    plasmacloud.cluster.generated.deployerClusterState = generatedDeployerClusterState;

    system.build.plasmacloudDeployerClusterState =
      jsonFormat.generate "plasmacloud-deployer-cluster-state.json" generatedDeployerClusterState;
  };
}