#!/usr/bin/python

### a utility to filter backups in the /db-backup area
#
# lists all files that match <prefix>-<date>-<time>.<suffix>
# and preserves only the following
# (*) more than 2 months old : one file per month
# (*) more than 2 weeks old : one file per week
# (*) recent stuff: unchanged
# 

import sys
import os, os.path
from datetime import datetime, timedelta
from glob import glob
import re
import traceback

from optparse import OptionParser

now=datetime.now()
counter=0

class FileIgnored (Exception): pass

# in days
LEAVE_FILES_YOUNGER_THAN = 20
# keep that amount of (plain) months organized in weeks
# should not exceed 11
KEEP_FILES_IN_MONTH_AFTER = 2

# monday=0 sunday=6
# we prefer saturday as this is the end of the week
PREFERRED_WEEKDAY=5

#
## utility to parse a filename
datetime_pattern="(?P<datetime>[0-9]{8}\-[0-9]{4})"
# returns a tuple
# (bool, prefix, suffix, datetime)

prefix_pattern="(?P<prefix>[\w\-]+)"
datetime_pattern="(?P<datetime>[0-9]{8}\-[0-9]{4})"
suffix_pattern="(?P<suffix>[\w\.]+)"
filename_matcher=re.compile("\A%s[-_]%s\.%s\Z"%(prefix_pattern,datetime_pattern,suffix_pattern))

parsing_failed = (False, None, None, None)

def parse_filename (filename):
    match=filename_matcher.match(filename)
    if not match:
        return parsing_failed
    else:
        d = match.groupdict()
        try:
            dt = datetime.strptime(d['datetime'],"%Y%m%d-%H%M")
            return (True, d['prefix'], d['suffix'], dt)
        except:
            print "failed to parse timestamp %s from %s"%(d['datetime'],filename)
            traceback.print_exc()
            return parsing_failed

# one entry like this per file, managed in the Kind class
class File:

    def __init__ (self, dir, filename, datetime, options):
        self.dir=dir
        self.filename=filename
        self.datetime=datetime
        self.options=options
        self.age=now-datetime
        self.weekday=self.datetime.weekday()
        if self.age.days<0:
            if self.options.verbose: print 'Filename %s is from the future - skipped'%filename
            raise FileIgnored,"Filename from the future %s"%filename
        self.group = self._group_string()

    def __repr__ (self):
        return "%s (%s) -- weekday %s"%(self.path(),self.datetime,self.datetime.weekday())

    def path (self):
        return os.path.normpath(os.path.join(self.dir,self.filename))

    def age_days (self):
        return self.age.days

    # oldest first
    @staticmethod
    def sort_age (file1, file2):
        # for 2.7, seems safer
        try:
            return int((file1.datetime-file2.datetime).total_seconds())
        # otherwise resort to days, we should be clear as our backups are daily
        except:
            return int((file1.datetime-file2.datetime).days)
        
    @staticmethod
    def sort_relevance (file1, file2):
        w1=file1.weekday
        w2=file2.weekday
        if w1==PREFERRED_WEEKDAY and w2==PREFERRED_WEEKDAY:
            return File.sort_age (file1,file2)
        elif w1!=PREFERRED_WEEKDAY and w2!=PREFERRED_WEEKDAY:
            return File.sort_age (file1,file2)
        elif w1==PREFERRED_WEEKDAY and w2!=PREFERRED_WEEKDAY:
            return -1
        elif w1!=PREFERRED_WEEKDAY and w2==PREFERRED_WEEKDAY:
            return 1
    
    month_barrier=None
    week_barrier=None
    @staticmethod
    def _compute_barriers ():
        if File.month_barrier:
            return
        # find the exact datetime for a month change 
        # KEEP_FILES_IN_MONTH_AFTER months ago +
        year=now.year
        month=now.month
        day=1
        if now.month>=KEEP_FILES_IN_MONTH_AFTER+1:
            month -= KEEP_FILES_IN_MONTH_AFTER
        else:
            year -= 1
            month += (12-KEEP_FILES_IN_MONTH_AFTER)
        File.month_barrier = datetime (year=year, month=month, day=day)
        # find the next monday morning
        remaining_days=(7-File.month_barrier.weekday())%7
        File.week_barrier=File.month_barrier+timedelta(days=remaining_days)

    @staticmethod
    def compute_month_barrier(): File._compute_barriers(); return File.month_barrier
    @staticmethod
    def compute_week_barrier(): File._compute_barriers(); return File.week_barrier

    # returns a key for grouping files, the cleanup then
    # preserving one entry in the set of files with same group
    def _group_string (self):
        if self.age.days<=LEAVE_FILES_YOUNGER_THAN:
            if self.options.verbose: print 'Filename %s is recent (%d d) - skipped'%\
                    (self.filename,LEAVE_FILES_YOUNGER_THAN)
            raise FileIgnored,"Filename %s is recent"%self.filename
        # in the month range
        if self.datetime <= File.compute_month_barrier():
            return self.datetime.strftime("%Y%m")
        else:
            weeks=(self.datetime-File.compute_week_barrier()).days/7
            weeks += 1
            return "week%02d"%weeks

    def cleanup (self, preserved):
        global counter
        counter+=1
        src = os.path.abspath(os.path.basename(self.path()));
        if self.options.destination:
            dst = os.path.abspath(self.options.destination) + '/' + os.path.basename(self.path())
            if self.options.verbose:
                print "moving %s\n\tto %s"%(self.path(), dst)
            if not self.options.dry_run:
                os.rename (src, dst)
        else:
            if self.options.verbose:
               print "Would cleanup %s"%(src)
               print "    (keeping %s)"%preserved.path()
            if not self.options.dry_run:
                if self.options.verbose: print "unlink",src
                os.unlink (src)

# all files in a given timeslot (either month or week)
class Group:
    def __init__ (self, groupname):
        self.groupname=groupname
        self.files=[]
        self.count = 0
    def insert (self, file):
        self.files.append(file)
    def epilogue (self):
        self.files.sort (File.sort_relevance)
    def keep_one (self):
        for file in self.files[1:]:
            file.cleanup(self.files[0])
            self.count += 1

# all files with the same (prefix, suffix)
class Kind:

    def __init__ (self, prefix, suffix, options):
        self.prefix=prefix
        self.suffix=suffix
        self.options=options
        self.todelete = 0
        self.oldest = None
        self.newest = None
        # will contain tuples (filename, datetime)
        self.list = []

    # allow for basic checking to be done in File
    def add_file (self, dir, filename, datetime):
        try:
            self.list.append ( File (dir, filename, datetime, self.options) )
            self.newest = datetime
            if not self.oldest:
                self.oldest = datetime
        except FileIgnored: pass
        except:
            print 'could not append %s'%filename
            traceback.print_exc()
            pass

    def epilogue (self):
        # sort self.list according to file age, oldest first
        self.list.sort(File.sort_age)
        # prepare groups according to age
        self.groups = {}
        for file in self.list:
            groupname=file.group
            if groupname not in self.groups:
                self.groups[groupname]=Group(groupname)
            self.groups[groupname].insert(file)
        for group in self.groups.values():
            group.epilogue()

    def show (self):
        if not self.options.verbose: return
        print 30*'-',"%s-<date>.%s"%(self.prefix,self.suffix)
        entries=len(self.list)
        print " %d entries" % entries,
        if entries >=1:
            f=self.list[0]
            print " << %s - %s d old"%(f.filename, f.age_days()),
        if entries >=2:
            f=self.list[-1]
            print "|| %s - %s d old >>"%(f.filename, f.age_days())
        groupnames=self.groups.keys()
        groupnames.sort()
        groupnames.reverse()
        if self.options.extra_verbose:
            print " Found %d groups"%len(groupnames)
            for g in groupnames: 
                print "  %s"%g
                files=self.groups[g].files
                for file in files:
                    print "    %s"%file
        elif self.options.verbose:
            print " Found %d groups"%len(groupnames),
            for g in groupnames: print "%s->%d"%(g,len(self.groups[g].files)),
            print ''

    # sort on number of entries
    @staticmethod
    def sort_size (k1, k2):
        return len(k1.list)-len(k2.list)

    def cleanup (self):
        groupnames=self.groups.keys()
        groupnames.sort()
        for groupname in groupnames:
            if self.options.extra_verbose: print 'GROUP',groupname
            self.groups[groupname].keep_one()
            self.todelete += self.groups[groupname].count

# keeps an index of all files found, index by (prefix, suffix), then sorted by time
class Index:
    def __init__ (self,options):
        self.options=options
        self.index = {}

    def insert (self, dir, filename, prefix, suffix, datetime):
        key= (prefix, suffix)
        if key not in self.index:
            self.index[key] = Kind(prefix,suffix, self.options)
        self.index[key].add_file (dir, filename, datetime)

    # we're done inserting, do housecleaning
    def epilogue (self):
        for (key, kind) in self.index.items():
            kind.epilogue()

    def show (self):
        # sort on number of entries
        kinds = self.index.values()
        kinds.sort (Kind.sort_size)
        for kind in kinds:
            kind.show()

    def insert_many (self, dir, filenames):
        for filename in filenames:
            (b,p,s,d) = parse_filename (filename)
            if not b:
		if self.options.verbose:
                    print "Filename %s does not match - skipped"%filename
                continue
            self.insert (dir, filename, p, s, d)

    def cleanup (self):
        for kind in self.index.values():
            kind.cleanup()
            
    def summary (self):
        print "%-30s%-10s%10s%25s%25s"%("Prefix","Suffix","Num (Del)","Oldest","Newest")
        for kind in self.index.values():
            print "%-30s%-10s%3s  (%3s) %30s%30s"%(kind.prefix, kind.suffix, len(kind.list), kind.todelete, kind.oldest, kind.newest)

def handle_dir_pattern (index, dir, pattern):
    try:
        os.chdir(dir)
    except:
        print "Cannot chdir into %s - skipped"%dir
        return
    filenames=glob(pattern)
    index.insert_many (dir, filenames)

def main ():
    usage="Usage: %prog [options] dir_or_files"
    parser=OptionParser(usage=usage)
    parser.add_option ("-v","--verbose",dest='verbose',action='store_true',default=False,
                       help="run in verbose mode")
    parser.add_option ("-x","--extra-verbose",dest='extra_verbose',action='store_true',default=False,
                       help="run in extra verbose mode")
    parser.add_option ("-n","--dry-run",dest='dry_run',action='store_true',default=False,
                       help="dry run")
    parser.add_option ("-m","--move-to",dest='destination',action='store',type='string',default=False,
                       help="move to <destination> instead of removing the file")
    parser.add_option ("-o","--offset",dest='offset',action='store',type='int',default=0,
                       help="pretend we run <offset> days in the future")
    parser.add_option ("-s","--summary",dest='summary',action='store_true',default=False,
                       help="print a summary")
    (options, args) = parser.parse_args()
    if options.extra_verbose: options.verbose=True
    try:
        if options.offset !=0:
            global now
            now += timedelta(days=options.offset)
    except:
        traceback.print_exc()
        print "Offset not understood %s - expect an int. number of days"%options.offset
        sys.exit(1)
    
    if options.destination and not os.path.isdir(options.destination):
        print "Destination should be a directory"
        sys.exit(1)

    if len(args) ==0:
        parser.print_help()
        sys.exit(1)

    # args can be directories, or patterns, like 
    # main /db-backup /db-backup-f8/*bz2
    # in any case we handle each arg completely separately
    dir_patterns=[]
    for arg in args:
        if os.path.isdir (arg):
            if arg=='.': arg=os.getcwd()
            dir_patterns.append ( (arg, '*',) )
        else:
            (dir,pattern)=os.path.split(arg)
            if not dir: dir=os.getcwd()
            dir_patterns.append ( (dir, pattern,) )

    index = Index (options)
    for (dir, pattern) in dir_patterns: handle_dir_pattern (index, dir, pattern)
    index.epilogue()
    index.show()
    index.cleanup()
    if (options.summary) :
        index.summary()
    if options.verbose:
        print 'Found %d entries to unlink'%counter
            
if __name__ == '__main__':
    main()