#!/usr/bin/tclsh9.0

set exitstatus 0
set puavo_btrfs_root "/.puavo"

if {[catch { exec -ignorestderr mountpoint -q $puavo_btrfs_root }]} {
  exit $exitstatus
}

proc puavo_conf key { exec -ignorestderr puavo-conf $key }

proc get_btrfs_size puavo_btrfs_root {
  set btrfs_output [
    exec -ignorestderr env LANG=C \
         btrfs filesystem usage --raw $puavo_btrfs_root
  ]

  foreach line [split $btrfs_output "\n"] {
    set columns [regexp -all -inline {\S+} $line]
    if {[lindex $columns 0] eq "Device" && [lindex $columns 1] == "size:"
      && [string is integer -strict [lindex $columns 2]]} {
        return [lindex $columns 2]
    }
  }

  error "could not determine btrfs size"
}

proc convert_limitspec_to_bytes limitspec {
  if {![regexp {^([0-9]+)(B|K|M|G|T|P)$} $limitspec _ value suffix]} {
    error "can not understand the quota limit"
  }
  set exp [lsearch { B K M G T P } $suffix]
  expr { int($value * 1024**$exp) }
}

proc get_limits_by_subvolume btrfs_size {
  global exitstatus

  set using_overlay [expr { [puavo_conf puavo.image.overlay] ne "" }]

  set puavo_conf_vars [list puavo.admin.btrfs.quotas.base]
  set puavo_hosttype [puavo_conf puavo.hosttype]
  if {$puavo_hosttype in {bootserver laptop}} {
    lappend puavo_conf_vars "puavo.admin.btrfs.quotas.${puavo_hosttype}"
  }
  lappend puavo_conf_vars \
          puavo.admin.btrfs.quotas.mode puavo.admin.btrfs.quotas.tweaks

  # get quota space from puavo-conf variables and put those to dict
  # (if duplicates per subvolume, take the last one)
  set puavo_quota_specs [
    concat {*}[
      lmap var $puavo_conf_vars { split [puavo_conf $var] }
    ]
  ]

  set free_subvolume ""
  set quota_limits [dict create]
  foreach subvolume_quota_spec $puavo_quota_specs {
    try {
      lassign [split $subvolume_quota_spec :] subvolume limitspec
      if {$limitspec eq "FREE"} {
        if {$free_subvolume ne ""} { error "free subvolume already set" }
        set free_subvolume $subvolume
        continue
      }

      set limit [convert_limitspec_to_bytes $limitspec]

      if {$subvolume eq "imageoverlays" && $limit == 0 && $using_overlay} {
        # If using image overlay and imageoverlays quota is set to zero
        # (the default), set it to 5G so it can actually be used.
        # This is somewhat orthogonal but trying to imageoverlay with quota
        # set to zero is not a good experience.
        set limit 5368709120
      }

      dict set quota_limits $subvolume $limit
    } on error errmsg {
      puts stderr [
        format {error handling quota limit "%s": %s} \
               $subvolume_quota_spec $errmsg
      ]
      set exitstatus 1
    }
  }

  if {$free_subvolume ne ""} {
    set btrfs_spare_size 1073741824
    set the_sum_of_limits [tcl::mathop::+ {*}[dict values $quota_limits]]
    set available_space [
      expr { $btrfs_size - $btrfs_spare_size - $the_sum_of_limits }
    ]
    if {$available_space < 0} {
      puts stderr [
        format {the btrfs filesystem is not big enough: %d bytes missing} \
               [expr { - $available_bytes }]
      ]
      set available_space 0
    }
    dict set quota_limits $free_subvolume $available_space
  }

  return $quota_limits
}

proc get_current_quotas_by_path puavo_btrfs_root {
  set btrfs_output [
    exec -ignorestderr env LANG=C btrfs qgroup show -re --raw \
      $puavo_btrfs_root 2>/dev/null
  ]

  set quotas [dict create]

  foreach line [split $btrfs_output "\n"] {
    set columns [regexp -all -inline {\S+} $line]
    lassign $columns id referenced exclusive max_referenced max_exclusive path
    if {$path eq "<toplevel>"} continue
    if {[regexp {^[0-9]+/[0-9]+$} $id]} {
      dict set quotas $path [
        dict create id $id max_referenced $max_referenced
      ]
    }
  }

  return $quotas
}

proc check_quotas {limits_by_subvolume quotas_by_path} {
  dict for {path quotainfo} $quotas_by_path {
    set current_limit [dict get $quotainfo max_referenced]
    set expected_limit [dict getwithdefault $limits_by_subvolume $path 0]
    if {$current_limit != $expected_limit} {
      return false
    }
  }

  return true
}

proc enable_quotas puavo_btrfs_root {
  # Go through "disable" in case quotas have some state we are not handling
  # but we are supposed to clear.
  exec -ignorestderr btrfs quota disable   $puavo_btrfs_root
  exec -ignorestderr btrfs quota enable    $puavo_btrfs_root
  exec -ignorestderr btrfs quota rescan -w $puavo_btrfs_root
}

set btrfs_size          [get_btrfs_size $puavo_btrfs_root]
set limits_by_subvolume [get_limits_by_subvolume $btrfs_size]

try {
  set quotas_by_path [get_current_quotas_by_path $puavo_btrfs_root]
} on error {} {
  # If we fail to get current quotas, perhaps quotas are not enabled yet?
  # Enable quotas and retry.
  enable_quotas $puavo_btrfs_root
  set quotas_by_path [get_current_quotas_by_path $puavo_btrfs_root]
}

if {[check_quotas $limits_by_subvolume $quotas_by_path]} {
  exit $exitstatus
}

puts "quotas are not as expected, re-applying quotas"
enable_quotas $puavo_btrfs_root
# re-read the quota information in case disable/enable changed the queue IDs
set quotas_by_path [get_current_quotas_by_path $puavo_btrfs_root]

dict for {path quotainfo} $quotas_by_path {
  set limit [dict getwithdefault $limits_by_subvolume $path 0]
  try {
    set qgroup_id [dict get $quotainfo id]
    puts "applying quota limit ${limit} bytes to ${path} (${qgroup_id})"
    exec -ignorestderr btrfs qgroup limit $limit $qgroup_id $puavo_btrfs_root
  } on error errmsg {
    error [format {error setting quota limits for "%s": %s} $subvolume $errmsg]
    set exitstatus 1
  }
}

exit $exitstatus
