#!/usr/bin/python

# A server version of hammie.py


"""Usage: %(program)s [-h] [-o <spambayes option>]

Where:
    -h
        show usage and exit
    -o
        set SpamBayes options
"""

import os

# Wish there was a prettier way to do this
os.chdir ( os.path.expanduser ( '~/.spambayes' ) )

import getopt
import sys
import traceback
import StringIO
import gobject
import dbus
import dbus.service
import dbus.mainloop.glib
import cPickle

from spambayes import hammie, Options
from spambayes import storage
from spambayes import mboxutils

GCONF_DIR = '/apps/evolution/mail/junk/spambayes'
GCONF_TRAIN_ON_ALL_MESSAGES = GCONF_DIR + '/train_on_all_messages'
GCONF_HAM_CUTOFF = GCONF_DIR + '/ham_cutoff'
GCONF_SPAM_CUTOFF = GCONF_DIR + '/spam_cutoff'
MESSAGE_DB = os.path.join ( os.path.expanduser ( '~/.spambayes' ), 'evolution_plugin.db' )

import gconf

program = sys.argv[0] # For usage(); referenced by docstring above

class DBusHammie ( dbus.service.Object, hammie.Hammie ):
    def __init__ ( self, bus, path, bayes ):
        dbus.service.Object.__init__ ( self, bus, path )
        hammie.Hammie.__init__ ( self, bayes )

        # seen messages map from Message-ID to int (spam = 1, ham = 0)
        if os.path.exists ( MESSAGE_DB ):
            fp = open ( MESSAGE_DB, 'rb' )
            self.seenMessages = cPickle.load(fp)
            fp.close ()
        else:
            self.seenMessages = {}
        
        self.gconf_client = gconf.client_get_default ()
        self.gconf_client.add_dir ( GCONF_DIR, gconf.CLIENT_PRELOAD_NONE )
        self.gconf_client.notify_add ( GCONF_DIR, self.gconfKeyChanged )
        
        self.gconfKeyChanged ( self.gconf_client, -1, self.gconf_client.get_entry(GCONF_TRAIN_ON_ALL_MESSAGES, "", True) )
        self.gconfKeyChanged ( self.gconf_client, -1, self.gconf_client.get_entry(GCONF_HAM_CUTOFF, "", True) )
        self.gconfKeyChanged ( self.gconf_client, -1, self.gconf_client.get_entry(GCONF_SPAM_CUTOFF, "", True) )

    def gconfKeyChanged ( self, client, contextId, entry, *extra ):
        if entry.get_key() == GCONF_TRAIN_ON_ALL_MESSAGES:
            self.trainOnAllMessages = client.get_bool(GCONF_TRAIN_ON_ALL_MESSAGES)
        
        elif entry.get_key() == GCONF_HAM_CUTOFF:
            Options.options['Categorization', 'ham_cutoff'] = client.get_float(GCONF_HAM_CUTOFF)
            
        elif entry.get_key() == GCONF_SPAM_CUTOFF:
            Options.options['Categorization', 'spam_cutoff'] = client.get_float(GCONF_SPAM_CUTOFF)
        
        else:
            print 'Unkown gconf entry changed: ', entry.get_key()
    
    @dbus.service.method ( 'net.halfdans.Spam', in_signature = 'ay', out_signature = 'ay', byte_arrays = True )
    def filter(self, msg):
        print '=' * 78
        print 'Filter'
        print '-' * 78
        print msg
        bin = hammie.Hammie.filter(self, msg)
        print '-' * 78
        print bin
        return bin
    
    @dbus.service.method ( 'net.halfdans.Spam', in_signature = 'ayb', out_signature = '', byte_arrays = True )
    def train ( self, msg, train_spam ):
        t = 'ham'
        if train_spam:
            t = 'spam'
        print '=' * 78
        print 'Train as', t
        print '-' * 78
        
        emailMsg = mboxutils.get_message(msg)
        print emailMsg.as_string()
        print '-' * 78
        
        if 'Message-ID' in emailMsg:
            messageId = emailMsg['Message-ID']
            if messageId in self.seenMessages:
                if self.seenMessages[messageId] == int(train_spam):
                    print 'Message already seen and trained as', t
                else:
                    print 'Message already seen, untraining before training.'
                    self.untrain ( msg, self.seenMessages[messageId] != 0 )
            
            hammie.Hammie.train(self, emailMsg, train_spam)
            self.seenMessages[messageId] = int(train_spam)
        else:
            print 'Message has no Message-ID, skipping.'
    
    @dbus.service.method ( 'net.halfdans.Spam', in_signature = 'ayb', out_signature = '', byte_arrays = True )
    def untrain ( self, msg, train_spam ):
        t = 'ham'
        if train_spam:
            t = 'spam'
        print '=' * 78
        print 'UnTrain as', t
        print '-' * 78
        print msg
        hammie.Hammie.untrain(self, msg, train_spam)
    
    @dbus.service.method ( 'net.halfdans.Spam', in_signature = 'ay', out_signature = 'd', byte_arrays = True )
    def score ( self, msg ):
        try:
            print '=' * 78
            s = hammie.Hammie.score(self, msg)
            print 'Score ->', s
            print '-' * 78
            print msg
            if self.trainOnAllMessages:
                if s < Options.options['Categorization', 'ham_cutoff']:
                    self.train ( msg, False )
                elif s > Options.options['Categorization', 'spam_cutoff']:
                    self.train ( msg, True )
            return s
        except Exception, e:
            print 'ERROR: Score:', e
            traceback.print_exc()
            raise
    
    @dbus.service.method ( 'net.halfdans.Spam', in_signature = 'ay', out_signature = 'a(sd)', byte_arrays = True )
    def clues ( self, msg ):
        print '=' * 78
        (score, clues) = hammie.Hammie.score(self, msg, evidence = True)
        print 'Clues'
        print '-' * 78
        print msg
        print '-'
        print hammie.Hammie.formatclues(self, clues, '\n')
        return clues
    
    @dbus.service.method ( 'net.halfdans.Spam', in_signature = '', out_signature = 'i' )
    def ping ( self ):
        print 'Ping'
        return 1

    @dbus.service.method ( 'net.halfdans.Spam', in_signature = '', out_signature = '' )
    def exit ( self ):
        global mainloop
        
        print 'Exit'
        self.store ()
        fp = open ( MESSAGE_DB, 'wb' )
        cPickle.dump ( self.seenMessages, fp )
        fp.close ()
        mainloop.quit ()

def usage(code, msg=''):
    """Print usage message and sys.exit(code)."""
    if msg:
        print >> sys.stderr, msg
        print >> sys.stderr
    print >> sys.stderr, __doc__
    sys.exit(code)

def main():
    global mainloop
    
    """Main program; parse options and go."""
    try:
        opts, args = getopt.getopt(sys.argv[1:], 'ho:')
    except getopt.error, msg:
        usage(2, msg)

    options = Options.options

    for opt, arg in opts:
        if opt == '-h':
            usage(0)
        elif opt == '-o':
            options.set_from_cmdline(arg, sys.stderr)
    
    dbname, usedb = storage.database_type(opts)
    bayes = storage.open_storage(dbname, usedb)
    
    loop = dbus.mainloop.glib.DBusGMainLoop ( set_as_default = True )
    bus = dbus.SessionBus( mainloop = loop )
    bus_name = dbus.service.BusName ( 'net.halfdans.Spam', bus )
    hammie = DBusHammie(bus, '/net/halfdans/SpamBayes', bayes)
    mainloop = gobject.MainLoop()
    mainloop.run ()

if __name__ == "__main__":
    main()
