#!/usr/bin/python2

import argparse, glob, operator, os, random, shutil, subprocess, sys, signal
shouldQuit = False

def getFreeSpaceForDir(dir):
    stats = os.statvfs(dir)
    return (stats.f_bavail * stats.f_frsize)

def getFreePercentForDir(dir):
    stats = os.statvfs(dir)
    total = (stats.f_blocks)
    avail = (stats.f_bavail)
    return (total - avail) / float(total)

def getFileSize(fullPath):
    return os.path.getsize(fullPath)

def sizeof_fmt(num, suffix='B'):
    for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
        if abs(num) < 1024.0:
            return "%3.1f %s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f %s%s" % (num, 'Yi', suffix)

def signal_handler(signal, frame):
    print "\nWill quit when file has been moved.\nMoving File..."
    global shouldQuit
    shouldQuit = True

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--checkonly', action='store_true', help="Check only, don't move any files.")
    parser.add_argument('-p', '--percent', type=int, default=7, help="The percentage difference between the most full dir and least full dir that will stop balancing.")
    cmdargs = parser.parse_args()

    SGDIRS = []
    SGgrp = "Default"

    signal.signal(signal.SIGINT, signal_handler)

    print "\nBalance MythTV Storage Group Directories\nPress Ctrl+C to quit"

    # Get Storage Groups from MythDB
    try:
        from MythTV import MythDB
        mythDB = MythDB()
        records = mythDB.getStorageGroup()
    except:
        print "Couldn't connect to MythTV database."
        sys.exit(1)

    # Get Storage Group directories
    for record in records:
        if record.groupname == SGgrp:
            dirname = record.dirname
            SGDIRS.append(dirname)

    # If there are less than 2 directories defined bail as we can't move anything
    if len(SGDIRS) < 2:
        print "There are less than 2 directories defined. Exiting."
        sys.exit(0)

    while not shouldQuit:
        SGDIRSdata = []
        print "\n------------------------------------------------"
        print "'" + SGgrp + "' Storage Group Directories - Percent Used:"
        # Get percent free and size free
        for directory in SGDIRS:
            # Check if SG path exists
            if not os.path.exists(directory):
                print "  " + directory + " - Not Mounted"
                continue
            freePcent = getFreePercentForDir(directory)
            freeSize = getFreeSpaceForDir(directory)
            SGDIRSdata.append([directory, freePcent, freeSize])
            print "  %s - %.2f%%" % (directory, freePcent * 100)

        # Sort data on percent free
        SGDIRSdata = sorted(SGDIRSdata, reverse=True, key=operator.itemgetter(1))
        #print SGDIRSdata

        # Check if SG has any ts, mpg or nuv files
        i=0
        for dir in SGDIRSdata:
            mostFull = SGDIRSdata[i]
            i=i+1
            if len(glob.glob1(mostFull[0],"*.ts")) or len(glob.glob1(mostFull[0],"*.mpg")) or len(glob.glob1(mostFull[0],"*.nuv")):
                break
            else:
                if i == 1:
                    print "------------------------------------------------"
                print "  " + mostFull[0] + " - NO files to move"

        leastFull = SGDIRSdata[-1]

        print "------------------------------------------------"
        print "Most Used Storage Group Directory with files to move: "
        print "  %s - %.2f%%" % (mostFull[0], mostFull[1] * 100)
        print "Least Used Storage Group Directory: "
        print "  %s - %.2f%%" % (leastFull[0], leastFull[1] * 100)

        # Check if mostFull and leastFull are within the percent var of each other
        if mostFull[1] - (float(cmdargs.percent) / 100) < leastFull[1]:
            print "\nThe most used and least used storage group directories are\nwithin " + str(cmdargs.percent) + "% used of each other. No files will be moved."
            sys.exit()

        # Get random file from most used dir
        fileToMove = random.choice([f for f in os.listdir(mostFull[0]) if f.endswith(".ts") or f.endswith(".mpg") or f.endswith(".nuv")])
        filePathToMove = mostFull[0] + "/" + fileToMove

        # Check that the file isn't too big for least used dir
        fileSize = getFileSize(filePathToMove)
        if (fileSize > getFreeSpaceForDir(leastFull[0])):
            # Too big to move
            print filePathToMove + " is too big to move to " + leastFull[0]
            sys.exit()

        print "------------------------------------------------"
        print "Move File:"
        print "  " + filePathToMove
        print "  Size: " + sizeof_fmt(os.path.getsize(filePathToMove))
        print "To:"
        print "  " + leastFull[0]
        print "  Available: " + sizeof_fmt(getFreeSpaceForDir(leastFull[0]))

        # Move file
        if cmdargs.checkonly:
            print "------------------------------------------------"
            print"Check Only option was used. No files were moved."
            shouldQuit = True
        else:
            print "------------------------------------------------"
            print "Checking System Status..."
            if subprocess.call(["/usr/bin/python2", "/usr/LH/bin/idle.py", "-s"]):
                print "  System is busy. The file will not be moved."
                sys.exit()
            print "Moving File..."
            try:
                shutil.move(filePathToMove, leastFull[0])
            # eg. src and dest are the same file
            except shutil.Error as e:
                print('Error: %s' % e)
            except IOError as e:
                print('Error: %s' % e.strerror)

            # Remove png files
            print "------------------------------------------------"
            print "Removing png Files:"
            pngFiles = glob.glob(filePathToMove + "*.png")
            for p in pngFiles:
                os.remove(p)
                print "  " + p