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

with lib;

let
  cfg = config.services.naxos.storage;
in {
  options.services.naxos.storage = {
    enable = mkEnableOption "NaxOS Declarative ZFS Storage Engine";

    arcMaxBytes = mkOption {
      type = types.nullOr types.ints.positive;
      default = 4294967296; # 4 GiB default
      description = "Maximum ZFS ARC cache size in bytes.";
    };

    autoScrub = {
      enable = mkOption {
        type = types.bool;
        default = true;
        description = "Enable automated regular scrubbing of ZFS pools.";
      };
      interval = mkOption {
        type = types.str;
        default = "monthly";
        description = "Interval for auto-scrub (e.g. monthly, weekly).";
      };
    };

    autoTrim = {
      enable = mkOption {
        type = types.bool;
        default = true;
        description = "Enable periodic TRIM for SSD/NVMe vdevs.";
      };
    };

    autoSnapshot = {
      enable = mkOption {
        type = types.bool;
        default = true;
        description = "Enable automatic snapshot retention policies.";
      };
      hourly = mkOption {
        type = types.int;
        default = 24;
        description = "Number of hourly snapshots to keep.";
      };
      daily = mkOption {
        type = types.int;
        default = 7;
        description = "Number of daily snapshots to keep.";
      };
      weekly = mkOption {
        type = types.int;
        default = 4;
        description = "Number of weekly snapshots to keep.";
      };
      monthly = mkOption {
        type = types.int;
        default = 12;
        description = "Number of monthly snapshots to keep.";
      };
    };

    importExistingPools = mkOption {
      type = types.listOf types.str;
      default = [];
      example = [ "tank" "datapool" ];
      description = "List of existing foreign ZFS pools (e.g. from TrueNAS or nixos-lukas) to safely import without formatting.";
    };

    pools = mkOption {
      type = types.attrsOf (types.submodule {
        options = {
          enable = mkOption {
            type = types.bool;
            default = true;
            description = "Whether to manage this pool.";
          };
          ashift = mkOption {
            type = types.int;
            default = 12;
            description = "ZFS vdev ashift alignment value (12 = 4K sectors).";
          };
          layout = mkOption {
            type = types.enum [ "stripe" "mirror" "raidz1" "raidz2" "raidz3" ];
            default = "mirror";
            description = "Vdev topology layout.";
          };
          devices = mkOption {
            type = types.listOf types.str;
            default = [];
            example = [ "/dev/disk/by-id/nvme-..." "/dev/disk/by-id/ata-..." ];
            description = "Member disk devices or partitions.";
          };
          datasets = mkOption {
            type = types.attrsOf (types.submodule {
              options = {
                mountpoint = mkOption {
                  type = types.nullOr types.str;
                  default = null;
                  description = "Mountpoint for this dataset. If null, uses default ZFS mount.";
                };
                compression = mkOption {
                  type = types.enum [ "on" "off" "lz4" "zstd" "zstd-fast" "gzip" ];
                  default = "lz4";
                  description = "ZFS compression algorithm.";
                };
                recordsize = mkOption {
                  type = types.str;
                  default = "128K";
                  description = "Record size for dataset (e.g. 1M for media, 16K for databases).";
                };
                quota = mkOption {
                  type = types.nullOr types.str;
                  default = null;
                  description = "Optional dataset quota (e.g. 500G).";
                };
                reservation = mkOption {
                  type = types.nullOr types.str;
                  default = null;
                  description = "Optional dataset reservation.";
                };
                owner = mkOption {
                  type = types.str;
                  default = "root";
                  description = "POSIX owner of mountpoint.";
                };
                group = mkOption {
                  type = types.str;
                  default = "root";
                  description = "POSIX group of mountpoint.";
                };
                mode = mkOption {
                  type = types.str;
                  default = "0755";
                  description = "POSIX directory mode permissions.";
                };
              };
            });
            default = {};
            description = "Sub-datasets belonging to this pool.";
          };
        };
      });
      default = {};
      description = "Declaratively defined ZFS pools and datasets.";
    };
  };

  config = mkIf cfg.enable {
    boot.supportedFilesystems = [ "zfs" ];
    boot.kernelPackages = pkgs.linuxPackages;

    # Kernel parameter for ZFS ARC Max limit
    boot.kernelParams = mkIf (cfg.arcMaxBytes != null) [
      "zfs.zfs_arc_max=${toString cfg.arcMaxBytes}"
    ];

    # Safe pool import list (includes both configured pools and imported existing pools)
    boot.zfs.extraPools = unique (
      (attrNames (filterAttrs (n: v: v.enable) cfg.pools)) ++
      cfg.importExistingPools
    );
    boot.zfs.forceImportRoot = false;

    # Automated ZFS services
    services.zfs = {
      autoScrub = {
        enable = cfg.autoScrub.enable;
        interval = cfg.autoScrub.interval;
      };
      trim = {
        enable = cfg.autoTrim.enable;
      };
      autoSnapshot = {
        enable = cfg.autoSnapshot.enable;
        hourly = cfg.autoSnapshot.hourly;
        daily = cfg.autoSnapshot.daily;
        weekly = cfg.autoSnapshot.weekly;
        monthly = cfg.autoSnapshot.monthly;
      };
    };

    # System tools for storage diagnosis and disk monitoring
    environment.systemPackages = with pkgs; [
      zfs
      smartmontools
      hdparm
      nvme-cli
      parted
      gptfdisk
      iotop
      ncdu
      sanoid
      syncoid
    ];

    # Generate systemd.tmpfiles rules for configured datasets with custom owners/modes
    systemd.tmpfiles.rules = flatten (mapAttrsToList (poolName: poolCfg:
      mapAttrsToList (dsName: dsCfg:
        mkIf (dsCfg.mountpoint != null)
          "d ${dsCfg.mountpoint} ${dsCfg.mode} ${dsCfg.owner} ${dsCfg.group} -"
      ) poolCfg.datasets
    ) cfg.pools);

    # Generate fileSystems definitions for datasets with explicit mountpoints
    fileSystems = foldl' (acc: pool:
      let
        poolName = pool.name;
        poolCfg = pool.value;
      in acc // (mapAttrs' (dsName: dsCfg:
        nameValuePair dsCfg.mountpoint {
          device = "${poolName}/${dsName}";
          fsType = "zfs";
          options = [ "nofail" ];
        }
      ) (filterAttrs (n: v: v.mountpoint != null) poolCfg.datasets))
    ) {} (mapAttrsToList (name: value: { inherit name value; }) (filterAttrs (n: v: v.enable) cfg.pools));
  };
}