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

with lib;

let
  cfg = config.roles.btrfs;
in
{
  options.roles.btrfs = {
    enable = mkEnableOption "Enable btrfs filesystem management";
    
    filesystems = mkOption {
      type = types.attrsOf (types.submodule {
        options = {
          # Filesystem-level maintenance options
          scrub = {
            enable = mkOption { 
              type = types.bool; 
              default = true;
              description = "Enable automatic scrubbing for this filesystem";
            };
            interval = mkOption { 
              type = types.str; 
              default = "weekly";
              description = "Scrub interval (systemd timer format)";
            };
          };
          deduplication = {
            enable = mkOption { 
              type = types.bool; 
              default = false;
              description = "Enable beesd deduplication for this filesystem";
            };
            hashTableSizeMB = mkOption { 
              type = types.int; 
              default = 1024;
              description = "Hash table size in MB (should be multiple of 16)";
            };
            verbosity = mkOption { 
              type = types.str; 
              default = "info";
              description = "Logging verbosity level";
            };
          };
          balance = {
            enable = mkOption { 
              type = types.bool; 
              default = false;
              description = "Enable periodic balance operations";
            };
            interval = mkOption { 
              type = types.str; 
              default = "monthly";
              description = "Balance interval (systemd timer format)";
            };
            dataUsage = mkOption { 
              type = types.int; 
              default = 50;
              description = "Data usage threshold for balance";
            };
            metadataUsage = mkOption { 
              type = types.int; 
              default = 50;
              description = "Metadata usage threshold for balance";
            };
          };
          
          # Mountpoint-based configuration
          mountpoints = mkOption {
            type = types.attrsOf (types.submodule {
              options = {
                subvolume = mkOption { 
                  type = types.nullOr types.str; 
                  default = null;
                  description = "Subvolume name. If null, uses default subvolume.";
                };
                compression = mkOption { 
                  type = types.str; 
                  default = "zstd";
                  description = "Compression algorithm (zstd, lzo, lz4, none)";
                };
                autodefrag = mkOption { 
                  type = types.bool; 
                  default = false;
                  description = "Enable automatic defragmentation";
                };
                extraOptions = mkOption { 
                  type = types.listOf types.str; 
                  default = [];
                  description = "Additional mount options";
                };
              };
            });
            default = {};
            description = "Mountpoint configurations for this filesystem";
          };
        };
      });
      default = {};
      description = "Btrfs filesystems configuration";
    };
  };

  config = mkIf cfg.enable {
    # Generate fileSystems configuration from mountpoints
    fileSystems = mkMerge (flatten (mapAttrsToList (device: fsCfg:
      mapAttrsToList (mountpoint: mountCfg:
        {
          ${mountpoint} = {
            device = device;
            fsType = "btrfs";
            options = 
              (optional (mountCfg.subvolume != null) "subvol=${mountCfg.subvolume}") ++
              [ "compress=${mountCfg.compression}" ] ++
              (optional mountCfg.autodefrag "autodefrag") ++
              mountCfg.extraOptions;
          };
        }
      ) fsCfg.mountpoints
    ) cfg.filesystems));
    
    # Configure scrub service using NixOS built-in
    services.btrfs.autoScrub = mkIf (any (fs: fs.scrub.enable) (attrValues cfg.filesystems)) {
      enable = true;
      interval = "weekly"; # TODO: Make this configurable per filesystem
      fileSystems = attrNames (filterAttrs (_: fs: fs.scrub.enable) cfg.filesystems);
    };
    
    # Configure beesd for filesystems with deduplication enabled
    services.beesd.filesystems = mapAttrs' (device: fsCfg: 
      nameValuePair (replaceStrings ["/"] ["_"] (replaceStrings ["-"] ["_"] device)) {
        spec = device;
        hashTableSizeMB = fsCfg.deduplication.hashTableSizeMB;
        verbosity = fsCfg.deduplication.verbosity;
      }
    ) (filterAttrs (_: fs: fs.deduplication.enable) cfg.filesystems);
    
    # Custom balance services for filesystems with balance enabled
    systemd.services = mkMerge (mapAttrsToList (device: fsCfg: mkIf fsCfg.balance.enable {
      "btrfs-balance-${replaceStrings ["/"] ["-"] (replaceStrings ["-"] ["_"] device)}" = {
        description = "Balance btrfs filesystem ${device}";
        script = ''
          ${pkgs.btrfs-progs}/bin/btrfs balance start \
            -dusage=${toString fsCfg.balance.dataUsage} \
            -musage=${toString fsCfg.balance.metadataUsage} \
            ${device}
        '';
        serviceConfig = {
          Type = "oneshot";
          Nice = 19;
          IOSchedulingClass = "idle";
        };
      };
    }) cfg.filesystems);
    
    # Balance timers
    systemd.timers = mkMerge (mapAttrsToList (device: fsCfg: mkIf fsCfg.balance.enable {
      "btrfs-balance-${replaceStrings ["/"] ["-"] (replaceStrings ["-"] ["_"] device)}" = {
        description = "Periodic balance for ${device}";
        wantedBy = [ "timers.target" ];
        timerConfig = {
          OnCalendar = fsCfg.balance.interval;
          Persistent = true;
        };
      };
    }) cfg.filesystems);
  };
}