#!/usr/bin/python3
#
# Copyright (C) 2016 Hans van Kranenburg <hans@knorrie.org>
#
# This file is part of python-btrfs.
#
# python-btrfs is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# python-btrfs is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with python-btrfs.  If not, see <http://www.gnu.org/licenses/>.


import btrfs
import sys


def munin_config(fs):
    print("multigraph btrfs_usage_{0}".format(str(fs.fsid).replace('-', '_')))
    print("graph_args --base 1024 -l 0")
    print("graph_vlabel bytes")
    print("graph_title btrfs space usage for {0}".format(fs.path))
    print("graph_category disk")
    print("graph_info This graph shows how btrfs uses available space")

    if not fs.mixed_groups():
        print("data_used.label Used Data")
        print("data_used.draw AREA")
        print("data_used.info Used Data")
        print("data_used.colour 33FF33")
        print("data_unused.label Unused Data")
        print("data_unused.draw STACK")
        print("data_unused.info Unused Data")
        print("data_unused.colour 00CC00")
        print("metadata_used.label Used Metadata")
        print("metadata_used.draw STACK")
        print("metadata_used.info Used Metadata")
        print("metadata_used.colour 3399FF")
        print("metadata_unused.label Unused Metadata")
        print("metadata_unused.draw STACK")
        print("metadata_unused.info Unused Metadata")
        print("metadata_unused.colour 0000CC")
    else:
        print("data_metadata_used.label Used Data+Metadata")
        print("data_metadata_used.draw AREA")
        print("data_metadata_used.info Used Data+Metadata")
        print("data_metadata_used.colour 99FFEE")
        print("data_metadata_unused.label Unused Data+Metadata")
        print("data_metadata_unused.draw STACK")
        print("data_metadata_unused.info Unused Data+Metadata")
        print("data_metadata_unused.colour 669999")

    print("system_used.label Used System")
    print("system_used.draw STACK")
    print("system_used.info Used System")
    print("system_used.colour FFFF33")
    print("system_unused.label Unused System")
    print("system_unused.draw STACK")
    print("system_unused.info Unused System")
    print("system_unused.colour CCCC00")

    print("parity.label Parity Blocks")
    print("parity.draw STACK")
    print("parity.info Parity Blocks")
    print("parity.colour 9900FF")

    print("unallocated.label Unallocated")
    print("unallocated.draw STACK")
    print("unallocated.info Not allocated raw space")
    print("unallocated.colour FFFFFF")
    print("non_alloc_reclaimable.label Reclaimable non-alloc")
    print("non_alloc_reclaimable.draw STACK")
    print("non_alloc_reclaimable.info Reclaimable not allocatable")
    print("non_alloc_reclaimable.colour BBBBBB")
    print("non_alloc.label Non-allocatable")
    print("non_alloc.draw STACK")
    print("non_alloc.info Non-allocatable")
    print("non_alloc.colour 888888")
    print("total.label Total")
    print("total.draw LINE2")
    print("total.info Total raw space")
    print("total.colour 000000")
    print("")


def munin_values(fs):
    print("multigraph btrfs_usage_{0}".format(str(fs.fsid).replace('-', '_')))

    # Get detailed usage statistics.
    usage = fs.usage()

    # Whatever happens, we should not stack the graph above this. IOW, the
    # unallocated bytes we end up with is just whatever is left over after
    # doing all other things.
    left = usage.total

    if not fs.mixed_groups():
        data_used = usage.block_group_type_usage[btrfs.BLOCK_GROUP_DATA].used
        data_allocated = usage.block_group_type_usage[btrfs.BLOCK_GROUP_DATA].allocated
        metadata_used = usage.block_group_type_usage[btrfs.BLOCK_GROUP_METADATA].used
        metadata_allocated = usage.block_group_type_usage[btrfs.BLOCK_GROUP_METADATA].allocated
        print("data_used.value {}".format(data_used))
        print("data_unused.value {}".format(data_allocated - data_used))
        print("metadata_used.value {}".format(metadata_used))
        print("metadata_unused.value {}".format(metadata_allocated - metadata_used))
        left = left - data_allocated - metadata_allocated
    else:
        mixed_type = btrfs.BLOCK_GROUP_DATA | btrfs.BLOCK_GROUP_METADATA
        used = usage.block_group_type_usage[mixed_type].used
        allocated = usage.block_group_type_usage[mixed_type].allocated
        print("data_metadata_used.value {}".format(used))
        print("data_metadata_unused.value {}".format(allocated - used))
        left -= allocated
    system_used = usage.block_group_type_usage[btrfs.BLOCK_GROUP_SYSTEM].used
    system_allocated = usage.block_group_type_usage[btrfs.BLOCK_GROUP_SYSTEM].allocated
    print("system_used.value {}".format(system_used))
    print("system_unused.value {}".format(system_allocated - system_used))
    left -= system_allocated

    print("parity.value {}".format(usage.parity))
    left -= usage.parity
    print("non_alloc_reclaimable.value {}".format(usage.unallocatable_reclaimable))
    left -= usage.unallocatable_reclaimable
    print("non_alloc.value {}".format(usage.unallocatable_hard))
    left = max(left - usage.unallocatable_hard, 0)
    print("unallocated.value {}".format(left))
    print("total.value {}".format(usage.total))
    print("")


def main():
    for path in btrfs.utils.mounted_filesystem_paths():
        with btrfs.FileSystem(path) as fs:
            if len(sys.argv) > 1 and sys.argv[1] == "config":
                munin_config(fs)
            else:
                munin_values(fs)


if __name__ == "__main__":
    main()
