#!/usr/bin/python
import sys, os, getopt
from ConfigParser import ConfigParser

# path in which this script lives. We assume hammie.py is in ..
if __name__=='__main__':
    mypath = sys.argv[0]
else:
    mypath = __file__
mypath = os.path.split(os.path.abspath(mypath))[0]

for sb_mboxtrain_path in [os.path.join(mypath,'..','scripts','sb_mboxtrain.py'),
                          os.path.join(mypath,'..','bin','sb_mboxtrain.py'),
                          '/usr/bin/sb_mboxtrain.py',
                          '/usr/local/bin/sb_mboxtrain.py',
                          os.path.join(mypath,'sb_mboxtrain.py') ]:
    if os.path.exists(sb_mboxtrain_path):
        break
else:
    sys.exit('cant find sb_mboxtrain.py')

class retrain:

    def __init__(self):
        pass
        
    def main(self):
        try:
            opts, args = getopt.getopt(sys.argv[1:], "ir", [])
        except getopt.GetoptError:
            # print help information and exit:
            sys.exit(self.usage())
        if len(args) != 0:
            sys.exit(self.usage())
        do = self.retrain
        for o, a in opts:
            if o == '-i':
                do = self.info
            elif o == '-r':
                do = self.retrain
            else:
                sys.exit(usage())
        do()
        
    def usage(self):
        return """
retrain.py [options]

options:
  -i    Display information about which folders are to be used
  -r    Perform retraining (default)
  -h    This help
"""
        
    def info(self):
        hams = []
        ignores = []
        print 'Spam folders:'
        for folder in self.folders():
            if self.is_spam(folder):
                print '    '+folder
            elif self.is_ignore(folder):
                ignores.append(folder)
            else:
                hams.append(folder)
        print '\nHam folders:'
        for folder in hams:
            print '    '+folder
        print '\nIgnored folders:'
        for folder in ignores:
            print '    '+folder
           
        
    def retrain(self):
        newdb = os.path.expanduser('~/.hammiedb.new')
        db = os.path.expanduser('~/.hammiedb')
        Mail = os.path.expanduser('~/Mail')
        if os.path.exists(newdb):
                os.unlink(newdb)
        cmd = [ sys.executable,
                sb_mboxtrain_path,
                '-d', newdb, '-f']
        for folder in self.folders():
            if self.is_spam(folder):
                cmd.append('-s')
                cmd.append(os.path.join(Mail,folder))
            elif self.is_ignore(folder):
                pass
            else:
                cmd.append('-g')
                cmd.append(os.path.join(Mail,folder))
        if os.spawnv(os.P_WAIT, cmd[0], cmd):
                sys.exit('not updating ~/.hammiedb')
        f = open(newdb)
        os.fsync(f.fileno()) # maybe unnecessary
        os.rename(newdb,db)
        
    def is_spam(self,folder):
        p = folder.lower().split('/')
        if 'spam' in p:
            return 1
        elif '.spam.directory' in p:
            return 1
        else:
            return 0

    def is_ignore(self,folder):
        return folder in ['sent-mail','drafts','outbox','trash','Last Search']

    def folders(self):
        s = []
        cp = ConfigParser()
        cp.read(os.path.expanduser('~/.kde/share/config/kmailrc'))
        for section in cp.sections():
            if section.startswith('Folder-') and len(section)>7:
                s.append(section[7:])
        return s

if __name__=='__main__':
    retrain().main()