#!/usr/bin/python
#
# Copyright (C) 2017 Hans van Kranenburg <hans@knorrie.org>
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import argparse
import btrfs
import errno
import heapq
import sys
import time


class Bork(Exception):
    pass


def load_block_groups(fs, max_used_pct):
    min_used_pct = 101
    block_groups = []
    print("Loading block group objects with used_pct <= {} ...".format(max_used_pct),
          end='', flush=True)
    for chunk in fs.chunks():
        if not (chunk.type & btrfs.BLOCK_GROUP_DATA):
            continue
        try:
            block_group = fs.block_group(chunk.vaddr, chunk.length)
            if block_group.used_pct <= max_used_pct:
                block_groups.append((block_group.used_pct, block_group))
            if block_group.used_pct < min_used_pct:
                min_used_pct = block_group.used_pct
        except IndexError:
            continue
    heapq.heapify(block_groups)
    print(" found {}".format(len(block_groups)))
    return min_used_pct, block_groups


def balance_one_block_group(fs, block_groups, max_used_pct):
    next_used_pct, next_block_group = block_groups[0]
    try:
        block_group = fs.block_group(next_block_group.vaddr, next_block_group.length)
    except IndexError:
        heapq.heappop(block_groups)
        return
    vaddr = block_group.vaddr
    used_pct = block_group.used_pct
    if used_pct > next_used_pct:
        if used_pct > max_used_pct:
            print("Ignoring block group vaddr {} used_pct changed {} -> {}".format(
                vaddr, next_used_pct, used_pct))
            heapq.heappop(block_groups)
        else:
            print("Postponing block group vaddr {} used_pct changed {} -> {}".format(
                vaddr, next_used_pct, used_pct))
            heapq.heapreplace(block_groups, (used_pct, block_group))
        return

    start_time = time.time()
    heapq.heappop(block_groups)
    args = btrfs.ioctl.BalanceArgs(vstart=vaddr, vend=vaddr+1)
    print("Balance block group vaddr {} used_pct {} ...".format(
        vaddr, used_pct), end='', flush=True)
    try:
        progress = btrfs.ioctl.balance_v2(fs.fd, data_args=args)
        end_time = time.time()
        print(" duration {} sec total {}".format(int(end_time - start_time), progress.considered))
    except KeyboardInterrupt:
        end_time = time.time()
        print(" duration {} sec".format(int(end_time - start_time)))
        raise


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-u',
        '--usage',
        type=int,
        default=100,
        help="Only consider block groups with less than this percentage used",
    )
    parser.add_argument(
        'mountpoint',
        help="Btrfs filesystem mountpoint",
    )
    return parser.parse_args()


def permission_check(fs):
    """This is a simple canary function that explodes if the user does not have
    enough permissions to use the search ioctl.
    """
    fs.top_level()


def main():
    args = parse_args()
    max_used_pct = args.usage
    path = args.mountpoint
    try:
        with btrfs.FileSystem(path) as fs:
            permission_check(fs)
            min_used_pct, block_groups = load_block_groups(fs, max_used_pct)
            if len(block_groups) == 0:
                print("Nothing to do, least used block group has used_pct {}".format(min_used_pct))
                return
            while len(block_groups) > 0:
                balance_one_block_group(fs, block_groups, max_used_pct)
    except OSError as e:
        if e.errno == errno.EPERM:
            raise Bork("Insufficient permissions to use the btrfs kernel API.\n"
                       "Hint: Try running the script as root user.".format(e))
        elif e.errno == errno.ENOTTY:
            raise Bork("Unable to retrieve data. Hint: Not a mounted btrfs file system?")
        raise


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print("Exiting...")
        sys.exit(130)  # 128 + SIGINT
    except Bork as e:
        print("Error: {0}".format(e), file=sys.stderr)
        sys.exit(1)
    except Exception as e:
        print(e)
        sys.exit(1)
