#!/usr/bin/env python

import sys,os,glob
from datetime import datetime
from datetime import timedelta
import sqlite3
import re,signal
import subprocess,threading

VERSION=2

W= '30'
R= '31'
G= '32'
Y= '33'
B= '34'
M= '35'
C= '36'
S= '1'
E= '0'
BR= '41'
CLR = '\033[2J'
SAVE = '\033[s'
LOAD = '\033[u'
CLRLN = '\033[K'
CLRBLN = '\033[1K'
DOWN = '\033[1B'

SORTCONVERT = lambda text: float(text) if is_number(text) else -1
SORTKEY = lambda key: SORTCONVERT(key[2][:-1])

def setup_options():
    ''' Setup the command line options '''
    from argparse import ArgumentParser
    import argparse

    parser=ArgumentParser(description='''
Tool to clean up and colorize the output of Anduril.
Example:   anduril run yourscript.and | %(prog)s 
You can tap in to an existing log with:
  tail -f -n +0 log/_global  | %(prog)s''',formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument("--no-colors",'--nc',action="store_false",dest="colors",default=True,
                        help="Disable colored output")
    parser.add_argument("-a",action="store_false",dest="activities",default=True,
                        help="Disable activities output")                    

    parser.add_argument("--version",action='version', version=VERSION)
    options=parser.parse_args()

    return options
def c(attribs):
    ''' ANSI colorizer '''
    if not options.colors:
        return ""
    return '\033['+';'.join(attribs)+'m'

def pos(y,x):
    ''' ANSI absolute position set '''
    return "\033["+str(y)+";"+str(x)+"H"


def colorize(string):
    ''' colorizes a string based on color_match '''
    if not options.colors:
        return string
    for c in color_match:
        string=color_match[c][0].sub(color_match[c][1],string)
    return string


def count_running(string, stats):
    ''' Counts the running executions '''
    
    spl=[i.strip() for i in string.split('|')]
    if len(spl)!=4:
        return stats
    spl.append(datetime.now())
    if spl[3] in stats['files']:
        index=stats['files'].index(spl[3])
        stats['running'][index]=spl
    else:
        stats['running'].append(spl)
    stats['running'].sort(key=SORTKEY, reverse=True)
    
    stats['files']=[i[3] for i in stats['running']]
    return stats

class EndProgram( Exception ):
    ''' Nice way of exiting the program '''
    pass

def remove_running(stats):
    ''' Remove Done files '''
    if not stats['running']:
        return stats
    # Remove Done/Fail older than 60sec
    for e in enumerate(stats['running']):
        if (datetime.now() - e[1][4] > timedelta(seconds=60)):
            if e[1][2]=='Done' or e[1][2]=='Failed':
                stats['running'].pop(e[0])
                stats['files'].pop(e[0])
    # Remove Done/Fail if there are too many:
    if len(stats['running'])>(stats['size'][0]-6):
        for e in enumerate(stats['running']):
            if e[1][2]=='Done':
                stats['running'].pop(e[0])
                stats['files'].pop(e[0])
                return stats
            if e[1][2]=='Failed':
                stats['running'].pop(e[0])
                stats['files'].pop(e[0])
                return stats    

    return stats
    
def is_number(s):
    ''' Check if string is float '''
    try:
        out=float(s)
        return True
    except:
        return False

def str_short(s,stats):
    ''' shorten text to fit screen '''
    maxL=stats['size'][1] - 16
    if len(s)<maxL:
        return s
    spl=s.split('/')
    sNew=spl[0]+'/...'+'/'.join(spl[1:])[-(maxL-len(spl[0])-5):]
    return sNew
    

def print_stats(stats):
    ''' Prints logged errors, and the status line '''
    #sys.stdout.write(SAVE)
    if options.activities:
        e=5
    else:
        e=0
    sys.stdout.write(pos(e+1,0)+c((S,C))+"=AeroFS Transfers"+"= "+c((E))+human_time()+
                     " DL: "+c((S,Y))+human_size(stats['pspeed'],0)+"/s "+c((E))+
                     "Cache: "+human_size(stats['psize'][4][0])+CLRLN)
    if (stats['running']):
        sys.stdout.write(pos(e+2,0)+"Last update: "+stats['running'][0][0]+CLRLN)
    else:
        return
    
    for ex in enumerate(stats['running']):
        sys.stdout.write(pos(e+3+ex[0],0)+'('+str(ex[0]+1).rjust(2)+') '+
                    ' '.join([ex[1][1],
                             colorize(ex[1][2].ljust(6)),
                             str_short(ex[1][3],stats)])+
                    CLRLN)
    for i in range(stats['size'][0]-6-len(stats['running'])):
        sys.stdout.write(pos(e+4+ex[0]+i,0)+"    "+CLRLN)
            
    sys.stdout.write(DOWN+CLRBLN+CLRLN)            
    #sys.stdout.write(LOAD)

def print_activities(activities):
    ''' Prints activity log '''
    if activities != None:
        activities.join()
        for e,l in enumerate(activities.readstdout().strip().split('\n')):
            sys.stdout.write(pos(e+1,0)+l+CLRLN)
    else:
        sys.stdout.write(pos(1,0)+'Recording activities..'+CLRLN)
    try:
        activities = Threaded("aerofs-sh activities -c 5")
        activities.start()
    except KeyboardInterrupt:
        pass
    except TypeError:
        activities.stop()
        activities.join()
    return activities
       
def human_time():
    return datetime.now().strftime("%I:%M:%S %p")

def human_size(size,precision=1):
    if size==None:
        return 'nan'
    suffixes=['B','KB','MB','GB','TB']
    suffixIndex = 0
    defPrecision=0
    while size > 1024:
        suffixIndex += 1
        size = size/1024.0
        defPrecision=precision
    return "%.*f%s"%(defPrecision,size,suffixes[suffixIndex])

def readinput(lf):
    try:
        line = lf.stdout.readline()
        #line=lf.readline()
        return line
    except:
        return "CleanerTimeout"

def termsize():
    rows, columns = os.popen('stty size', 'r').read().split()
    return (int(rows),int(columns))

def get_partial_dir():
    sql_file=os.path.join(os.path.expanduser("~"),
                             ".aerofs", "conf")
    conn=sqlite3.connect(sql_file)
    db=conn.cursor()
    conn.text_factory=str
    db.execute("SELECT v FROM c WHERE k = 'root'")
    for row in db:
        continue
    conn.close()
    cache_dir=glob.glob(os.path.join(os.path.dirname(row[0]),'.aerofs.aux*'))
    return os.path.join(cache_dir[0],'p')

def get_partial_size(dir):
    return sum([os.path.getsize(os.path.join(dir,f)) for f in os.listdir(dir) if os.path.isfile(os.path.join(dir,f))])

def partial_update(stats):
    ''' Calculate average speed of transfer '''
    stats['psize'].pop(0)
    stats['psize'].append( ( get_partial_size(stats['pdir']), datetime.now()) )
    speedlist=[]
    for i in range(len(stats['psize'])-1):
        sizediff=stats['psize'][i+1][0] - stats['psize'][i][0]
        timediff=(stats['psize'][i+1][1] - stats['psize'][i][1]).total_seconds()
        if timediff>0 and sizediff>0:
            speedlist.append( sizediff / timediff )
    if len(speedlist)==0:
        speed=0.0
    else:
        speed=sum(speedlist)/len(speedlist)
    
    stats['pspeed']=speed
    
    return stats

class Threaded(threading.Thread):
    def __init__(self,command):
        self.stdout = None
        self.stderr = None
        self.command=command
        self.p = None
        threading.Thread.__init__(self)

    def run(self):
        self.p = subprocess.Popen(self.command.split(),
                             shell=False,
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE)

    def readline(self):
        try:
            line = self.p.stdout.readline()
            #line=lf.readline()
            return line
        except:
            return "CleanerTimeout"
    
    def readstdout(self):
        self.stdout, self.stderr = self.p.communicate()       
        return self.stdout
    
    def stop(self):
        self.p.terminate()

options=setup_options()

color_match={#'line_ends':(re.compile('$'),c.END),
        'err':(re.compile('(Failed)'),c([R,S])+'\\1'+c([E])),
        'done':(re.compile('(Done)'),c([G,S])+'\\1'+c([E])),
        'percent':(re.compile('([0-9]+%)'),c([Y,S])+'\\1'+c([E])),
}

stats={'time':datetime.now()-timedelta(seconds=25), 
       'running':[],
       'files':[],
       'size': termsize(),
       'pdir': get_partial_dir(), 
       'psize': [( get_partial_size(get_partial_dir()), datetime.now())]*5,
       'pspeed': 0.0
       }

sys.stdout.write(CLR+pos(0,0)+"Launching...")
#proc = subprocess.Popen(['aerofs-sh','transfers'],stdout=subprocess.PIPE)
transfers = Threaded("aerofs-sh transfers")
transfers.start()

for e in range(5):
    sys.stdout.write(pos(e+1,0)+CLRLN)
activities = print_activities(None)

while 1:

    try:
        sys.stdout.flush()
        # set a 5 second timeout for the line read.
        signal.signal(signal.SIGALRM, transfers.readline)
        signal.alarm(5)
        line=transfers.readline()
        if not line:
            raise EndProgram
        
        if ( datetime.now() - stats['time']  > timedelta(seconds=30) ) and options.activities:
            activities=print_activities(activities)
            stats=partial_update(stats)
            stats['time'] = datetime.now()

        stats=remove_running(stats) 
        stats=count_running(line,stats)       

        print_stats(stats)
        
    except EndProgram,KeyboardInterrupt:
        transfers.stop()
        transfers.join()
        sys.stdout.write(DOWN+'\n')
        sys.stdout.flush()
        
        sys.exit(0)