import sys
import os
import re


def alert(arg):
    """Send an alert with a specific argument."""
    sys.stderr.write("%d%d%d\n" % (arg, arg, arg))


def humanbytes(B):
    """Return the given bytes as a human-friendly KB, MB, GB, or TB string."""
    B = float(B)
    KB = 1024
    MB = KB ** 2
    GB = KB ** 3
    TB = KB ** 4

    if B < KB:
        return '{0} {1}'.format(B, 'Bytes' if B != 1 else 'Byte')
    elif B < MB:
        return '{0:.2f} KB'.format(B / KB)
    elif B < GB:
        return '{0:.2f} MB'.format(B / MB)
    elif B < TB:
        return '{0:.2f} GB'.format(B / GB)
    else:
        return '{0:.2f} TB'.format(B / TB)


p = 0  # This constant will contain the result, transferred into the alert return value

try:
    # Detect local drives using `wmic` and regex
    detected_drives = os.popen('wmic logicaldisk get description,name | findstr "Local"').read()
    local_drives = re.findall(r'[A-Z]:', detected_drives)

    if not local_drives:
        print "No local drives detected."
        alert(0)
        sys.exit(0)

    for j in local_drives:
        cmd = os.popen('fsutil volume diskfree ' + j).read()
        print "Debug - Output of `fsutil` for drive {}:\n{}".format(j, cmd)  # Debugging the command output

        # Updated regex to handle commas
        total_space_match = re.search(r'Total\s+bytes\s+:\s+([\d,]+)', cmd)
        free_space_match = re.search(r'Total\s+free\s+bytes\s+:\s+([\d,]+)', cmd)

        if not total_space_match or not free_space_match:
            print "Failed to parse disk space for drive:", j
            continue

        # Remove commas from the matched values and convert to integers
        total_space = int(total_space_match.group(1).replace(',', ''))
        free_space = int(free_space_match.group(1).replace(',', ''))

        threshold = 5 * (total_space / 100)  # 5% threshold

        print "Total size of the local disk {}: {}".format(j, humanbytes(total_space))
        print "Free space in the local disk {}: {}".format(j, humanbytes(free_space))

        if free_space <= threshold:
            print "Low storage alert in local disk {}.".format(j)
            p = 1

except Exception as e:
    print "An error occurred:", str(e)
    alert(1)
    sys.exit(1)

# Send alert based on the flag
if p:
    alert(1)
else:
    alert(0)