#!/usr/bin/env python3

import sys
import re
import subprocess
import math
from optparse import OptionParser

###############################################################################
# Name: dump_attributes
# Args: path to directory
# Desc: executes getfattr to dump ceph attributes and returns as dictionary
###############################################################################
def dump_attributes(path):
	attrs = {}
	try:
		child = subprocess.Popen(["getfattr", "-d", "-m", "-", path],
			stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
	except OSError:
		print("Error executing getfattr. Is xattr installed?")
		sys.exit(1)
	if child.wait() != 0:
		print("Error executing getfattr. Is xattr installed?")
		sys.exit(1)
	output_string = child.stdout.read()
	fields = re.findall(r"^ceph\.dir\.([^=]+)=\"([^\"]+)\"$", output_string, re.MULTILINE)
	if len(fields) == 0:
		print(f'No ceph xattrs, is {path} in a ceph filesystem?')
		sys.exit(1)
	for field in fields:
		if field[0] != "rctime":
			attrs[field[0]] = field[1]
	return attrs

###############################################################################
# Name: display_attributes
# Args: path to directory
# Desc: calls dump_attributes and prints output
###############################################################################
def display_attributes(path):
	print(path, ":")
	attrs = dump_attributes(path)
	if "rbytes" in attrs.keys():
		attrs["rbytes"] = format_bytes(int(attrs["rbytes"]))
	max_width = len(max(attrs.values(), key=lambda x: len(x.split(" ")[0])).split(" ")[0])
	print("Files:                 ", "{0:>{1}}".format(attrs["files"], max_width) if "files" in attrs.keys() else "?")
	print("Directories:           ", "{0:>{1}}".format(attrs["subdirs"], max_width) if "subdirs" in attrs.keys() else "?")
	print("Recursive Files:       ", "{0:>{1}}".format(attrs["rfiles"], max_width) if "rfiles" in attrs.keys() else "?")
	print("Recursive Directories: ", "{0:>{1}}".format(attrs["rsubdirs"], max_width) if "rsubdirs" in attrs.keys() else "?")
	print("Total Size:            ", "{0:>{1}}".format(attrs["rbytes"].split(" ")[0], max_width) + " " + attrs["rbytes"].split(" ")[1] if "rbytes" in attrs.keys() else "?")
	print()

################################################################################
# Name: format_bytes
# Args: integer value in bytes
# Desc: formats size_bytes in SI base units and returns as string
################################################################################
def format_bytes(size_bytes):
	if size_bytes == 0:
		return "0 B"
	size_name = ("B", "KiB", "MiB", "GiB", "TiB", "PiB", "EiB", "ZiB", "YiB")
	i = int(math.floor(math.log(size_bytes, 1024)))
	p = math.pow(1024, i)
	s = round(size_bytes / p, 2)
	return "%s %s" % (s, size_name[i])

###############################################################################
# Name: main (cephfs-dir-stats)
# Args: (see parser)
# Desc: lists recursive ceph stats of specified directory
###############################################################################
def main():
	parser = OptionParser()
	(options, args) = parser.parse_args()
	if len(args) == 0:
		args = ["."]
	for arg in args:
		display_attributes(arg)

if __name__ == "__main__":
	main()
