#!/usr/bin/env python3

# Copyright (c) 2016, Antonio SJ Musumeci <trapexit@spawn.link>
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import argparse
import os
import subprocess
import sys


def ismergerfs(path):
    try:
        os.getxattr(path,'user.mergerfs.version')
        return True
    except IOError as e:
        return False


def mergerfs_control_file(basedir):
    return os.path.join(basedir,'.mergerfs')


def mergerfs_srcmounts(ctrlfile):
    srcmounts = os.getxattr(ctrlfile,'user.mergerfs.srcmounts')
    srcmounts = srcmounts.decode().split(':')
    return srcmounts


def find_a_file(src):
    for (dirpath,dirnames,filenames) in os.walk(src):
        if filenames:
            relfile = os.path.join(dirpath,filenames[0])
            relfile = relfile[len(src)+1:]
            return relfile


def execute(args):
    return subprocess.call(args)


def move_file(src,dst,relfile):
    frompath = os.path.join(src,'./',relfile)
    topath   = dst+'/'
    args = ['rsync',
            '-alXA',
            '--relative',
            '--progress',
            '--remove-source-files',
            frompath,
            topath]
    return execute(args)


def freespace_percentage(srcmounts):
    lfsp = []
    for srcmount in srcmounts:
        vfs = os.statvfs(srcmount)
        avail = vfs.f_bavail * vfs.f_frsize
        total = vfs.f_blocks * vfs.f_frsize
        per = avail / total
        lfsp.append((srcmount,per))
    return sorted(lfsp, key=lambda x: x[1])


def all_within_range(l,n):
    if len(l) == 0 or len(l) == 1:
        return True
    return (abs(l[0][1] - l[-1][1]) <= n)


def buildargparser():
    parser = argparse.ArgumentParser(description='balance files on a mergerfs mount based on percentage drive filled')
    parser.add_argument('dir',
                        type=str,
                        help='starting directory')
    parser.add_argument('-p',
                        dest='percentage',
                        type=float,
                        default=10.0,
                        help='percentage range of freespace (default 10.0)')
    return parser


def main():
    parser = buildargparser()
    args = parser.parse_args()

    args.dir = os.path.realpath(args.dir)

    ctrlfile = mergerfs_control_file(args.dir)
    if not ismergerfs(ctrlfile):
        print("%s is not a mergerfs mount" % args.dir)
        sys.exit(1)

    srcmounts = mergerfs_srcmounts(ctrlfile)

    percentage = args.percentage / 100
    try:
        l = freespace_percentage(srcmounts)
        while not all_within_range(l,percentage):
            fromdrive   = l[0][0]
            todrive     = l[-1][0]
            relfilepath = find_a_file(fromdrive)
            if not relfilepath:
                break
            rv = move_file(fromdrive,todrive,relfilepath)
            if rv:
                break;
            l = freespace_percentage(srcmounts)
    except KeyboardInterrupt:
        print("exiting: CTRL-C pressed")

    sys.exit(0)


if __name__ == "__main__":
   main()
