#!/usr/bin/python

import argparse, glob, operator, os, random, shutil, subprocess, sys, signal, time
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):
    if isBusy:
        sys.exit(0)
    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.")
    cmdargs = parser.parse_args()

    SGs = []
    SGDIRS = []
    SGgrp = "Default"
    SGselectdata = []

    signal.signal(signal.SIGINT, signal_handler)

    print("\nEmpty a MythTV Storage Group Directory\nPress Ctrl+C to quit")

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

    # get list of non duplicate SGs
    for record in records:
        SGs.append(record.groupname)
        SGs=list(set(SGs))
    # Move Default to top of list
    if "Default" in SGs:
        SGs.remove("Default")
        SGs.insert(0,"Default")

    # ask user which SG to use
    print("\n------------------------------------------------")
    print("Storage Groups:")
    for i,sg in enumerate(SGs):
        print(str(i+1) + ": " + sg)

    try:
        SGselect=input("\nEnter the number of the storage group to use (default 1): ") or 1
        SGselect=int(SGselect)
        if SGselect > len(SGs) or SGselect < 1:
            SGselect=int("e")
    except ValueError:
        print("You must enter a number between 1 and " + str(len(SGs)) + ". Exiting.")
        sys.exit(0)

    SGgrp=SGs[SGselect-1]

    # Get Storage Group directories
    for record in recs:
        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:")
        SGcnt=0
        # 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
            # Check if SG has data files to move
            if len(glob.glob(directory + "*.ts")) or len(glob.glob(directory + "*.mpg")) or len(glob.glob(directory + "*.nuv")) or len(glob.glob(directory + "*.jpg")):
                freePcent = getFreePercentForDir(directory)
                freeSize = getFreeSpaceForDir(directory)
                SGDIRSdata.append([directory, freePcent, freeSize])
                SGcnt=SGcnt+1
                print("%s: %s - %.2f%%" % (SGcnt, directory, freePcent * 100))
            else:
                # Check if the selected SG dir has no data files exit
                if SGselectdata and SGselectdata[0] == directory:
                    print("\n'" + SGgrp + "' Storage Group directories have no files to move. Exiting")
                    sys.exit(0)

        # Exit if no SGs with data found
        if SGcnt == 0:
            print("\n'" + SGgrp + "' Storage Group directories have no files to move. Exiting.")
            sys.exit(0)

        # Ask user to select which SG to empty if not already selected
        if not SGselectdata:
            try:
                SGDIRselect=int(input("\nEnter the number of the storage group directory to empty: "))
                if SGDIRselect > SGcnt or SGDIRselect < 1:
                    SGDIRselect=int("e")
            except ValueError:
                print("You must enter a number between 1 and %s. Exiting." %SGcnt)
                sys.exit(0)

            SGselectdata=SGDIRSdata[SGDIRselect-1]

        # Sort data on percent free
        SGDIRSdata = sorted(SGDIRSdata, reverse=True, key=operator.itemgetter(1))
        leastFull = SGDIRSdata[-1]

        # Make sure leastFull and SGselectdata are not the same dir
        if leastFull[0] == SGselectdata[0]:
            leastFull = SGDIRSdata[-2]

        # Get random file from user selected dir
        fileToMove = random.choice([f for f in os.listdir(SGselectdata[0]) if f.endswith(".ts") or f.endswith(".mpg") or f.endswith(".nuv") or f.endswith(".jpg")])
        filePathToMove = SGselectdata[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:
            global isBusy
            isBusy = True
            while isBusy:
                if shouldQuit:
                    sys.exit(0)
                print("------------------------------------------------")
                print("Checking System Status...")
                if subprocess.call(["/usr/bin/python", "/usr/bin/idle.py", "-s"]):
                    print("  System is busy. The file will not be moved.")
                    print("  Waiting 5 minutes before trying again.")
                    time.sleep(300)
                else:
                    isBusy = False
            print("Moving File...")
            try:
                shutil.move(filePathToMove, leastFull[0])
            # eg. src and dest are the same file
            except shutil.Error as e:
                a=input("\n%s. Overwrite destination (y/n)? " % e)
                if a == "y" or a == "Y":
                    os.remove(leastFull[0] + "/" + fileToMove)
                    shutil.move(filePathToMove, leastFull[0])
                else:
                    b=input("\nRemove %s (y/n)? " % filePathToMove)
                    if b == "y" or b == "Y":
                        os.remove(filePathToMove)
            # eg. source or destination doesn't exist
            except IOError as e:
                print(('Error: %s' % e.strerror))

            # Remove png files for Default & LiveTV SGs
            if SGgrp == "Default" or SGgrp == "LiveTV":
                print("------------------------------------------------")
                print("Removing png Files:")
                pngFiles = glob.glob(filePathToMove + "*.png")
                for p in pngFiles:
                    os.remove(p)
                    print("  " + p)
