#!/usr/bin/python3

#
# Imports
#
import sys
import time
import subprocess
import argparse


#sys.exit(1)

#
# Misc
#
# sys.tracebacklimit = 0


#
# Global variables
#
size_data_total = 0
size_data_exclusive = 0
size_snapshot_total = 0
size_snapshot_exclusive = 0


#
# Methods
#
def get_subvol_list(path):
    command = "sudo btrfs subvolume list -t %s" % (path)
    status, output = subprocess.getstatusoutput(command)

    if status != 0:
        raise Exception(command)

    # Every line contains the following values: subvol_id, gen, toplevel, path
    return output.splitlines()[2:]


def get_filesystem_size(path):
    command = "sudo btrfs filesystem show --raw %s" % (path)
    status, output = subprocess.getstatusoutput(command)

    if status != 0:
        # This command fails when running inside Docker container
        # return maximum size of any filesystem instead
        command = "sudo btrfs filesystem show --raw"
        status, output = subprocess.getstatusoutput(command)
        lines = output.splitlines()
        lines = [x for x in lines if "devid" in x]
        sizes = [int(line.split()[3]) for line in lines]
        return max(sizes)

    # The sizes are on the third line
    line = output.splitlines()[2]

    # Element 3 and 5 respectively contain total and used sizes
    return int(line.split()[3])


def get_id_root(name, path):
    lines = get_subvol_list(path)

    # Filter lines where toplevel == 5
    subvol_ids = [x for x in lines if int(x.split()[2]) == 5]

    # Try to retrieve the subvol_id for the root subvolume (if any)
    if len(subvol_ids) == 1:
        # The path contains a btrfs filesystem without subvolume for data
        return int(subvol_ids[0].split()[0])
    else:
        # The path contains a btrfs filesystem with multiple subvolumes for data
        try:
            return int(list(filter(lambda x: x.split()[3] == name, subvol_ids))[0].split()[0])
        except IndexError:
            pass

    # Volume not found, root is probably the btrfs default (5)
    return 5


def get_id_subvolumes(path, subvol_id):
    lines = get_subvol_list(path)
    lines = [x for x in lines if int(x.split()[2]) == subvol_id]
    return list([int(x.split()[0]) for x in lines])


def get_disk_usage(name, path):
    id_root = get_id_root(name, path)
    id_subvolumes = get_id_subvolumes(path, id_root)
    size_filesystem = get_filesystem_size(path)

    # Get disk usage from quota
    command = "sudo btrfs qgroup show --raw %s" % (path)
    status, output = subprocess.getstatusoutput(command)

    if status != 0:
        raise Exception(command)

    lines = output.splitlines()[2:]

    # Global variables
    global size_data_total
    global size_data_exclusive
    global size_snapshot_total
    global size_snapshot_exclusive

    # Total data volume in subvolume
    size_data_total = 0

    # Total data volume in snapshots
    #   -> this variable is useless
    size_snapshot_total = 0

    # Data exclusively in subvolume
    #   -> data that is not (yet) incorporated in a snapshot
    size_data_exclusive = 0

    # Data exclusively available in snapshots
    #   -> data that was removed from volume
    size_snapshot_exclusive = 0

    for line in lines:
        split = line.split()
        subvol_id = 0
        size_total = 0
        size_exclusive = 0
        try:
            subvol_id = int(split[0].split("/")[1])
            size_total = float(split[1])
            size_exclusive = float(split[2])
        except IndexError:
            # ignore "WARNING: Quota disabled"
            pass

        # size_exclusive is incorrect when snapshot is
        # removed and qgroups are not updated yet,
        # ignore the value when it seems unrealistic
        if size_exclusive > size_filesystem:
            size_exclusive = 0

        if subvol_id == id_root:
            size_data_total = size_total
            size_data_exclusive = size_exclusive
        elif subvol_id in id_subvolumes:
            size_snapshot_total += size_total
            size_snapshot_exclusive += size_exclusive


def rescan_quota(path):
    command = "sudo btrfs quota rescan %s" % (path)
    status, output = subprocess.getstatusoutput(command)
    if status != 0:
        Exception(command)


def print_human_readable(name):
    global size_data_total
    global size_data_exclusive
    global size_snapshot_exclusive
    size_data_total = size_data_total / (1024 * 1e6)
    size_data_exclusive = size_data_exclusive / (1024 * 1e6)
    size_snapshot_exclusive = size_snapshot_exclusive / (1024 * 1e6)
    print(
        "%10s: %6.1f Gb, %6.1f Gb, %6.1f Gb"
        % (name, size_data_total, size_data_exclusive, size_snapshot_exclusive)
    )


def print_rrd(name):
    timestamp = int(time.time())
    print(
        (
            "PUTVAL {}/exec-btrfs_{}/gauge-data_total {}:{:.1f}".format(
                hostname, name, timestamp, size_data_total
            )
        )
    )
    print(
        (
            "PUTVAL {}/exec-btrfs_{}/gauge-data_exclusive {}:{:.1f}".format(
                hostname, name, timestamp, size_data_exclusive
            )
        )
    )
    print(
        (
            "PUTVAL {}/exec-btrfs_{}/gauge-snapshot_total {}:{:.1f}".format(
                hostname, name, timestamp, size_snapshot_total
            )
        )
    )
    print(
        (
            "PUTVAL {}/exec-btrfs_{}/gauge-snapshot_exclusive {}:{:.1f}".format(
                hostname, name, timestamp, size_snapshot_exclusive
            )
        )
    )


#
# Volumes to scan
#
hostname = "shuttle"
interval = 10
volumes = list()

# SSD
volumes.append(["scratch",   "/host/root/mnt/mezzo/scratch"])
volumes.append(["sync",      "/host/root/mnt/mezzo/sync"])

# HDD
volumes.append(["personal",  "/host/root/mnt/yotta/helium/personal"])
volumes.append(["shared",    "/host/root/mnt/yotta/helium/shared"])
volumes.append(["neon",      "/host/root/mnt/yotta/neon"])
volumes.append(["krypton",   "/host/root/mnt/yotta/krypton"])
volumes.append(["borg",      "/host/root/mnt/yotta/xenon/borg"])
volumes.append(["rsnapshot", "/host/root/mnt/yotta/xenon/rsnapshot"])

#
# Command line arguments
#
parser = argparse.ArgumentParser(description="Get BTRFS disk usage")
parser.add_argument("-s", action="store_true", help="print in human readable format")
args = parser.parse_args()
human_readable = args.s


#
# Main
#
if human_readable:
    for (name, path) in volumes:
        get_disk_usage(name, path)
        print_human_readable(name)
else:
    # RRD mode
    while True:
        for (name, path) in volumes:
            get_disk_usage(name, path)

            print_rrd(name)

            sys.stdout.flush()
        time.sleep(interval)
        # rescan_quota(path)
