'''Accounting module for MediaProxy

Copyright 2004-2006 AG Projects
'''

import threading
import Queue
import sys
import os
import atexit
import cPickle
import traceback

from utilities import *
from serialize import SafeSerializer, SerializeError, UnserializeError
from configuration import readSettings, ConfigSection, dumpSection, Boolean
from radius import *
from mysql import *

class AccountingConfig(ConfigSection):
    accounting = 'none'

class DatabaseConfig(ConfigSection):
    table            = 'radacct'
    startTimeField   = 'AcctStartTime'
    stopTimeField    = 'AcctStopTime'
    durationField    = 'AcctSessionTime'
    callerBytesField = 'AcctInputOctets'
    calledBytesField = 'AcctOutputOctets'
    sessionIdField   = 'AcctSessionId'
    userAgentField   = 'SipUserAgents'
    codecTypesField  = 'SipCodecs'
    streamTypeField  = 'SipApplicationType'
    fromTagField     = 'SipFromTag'
    toTagField       = 'SipToTag'
    startInfoField   = 'ConnectInfo_start'
    stopInfoField    = 'ConnectInfo_stop'
    mediaInfoField   = 'MediaInfo'
    normalizedField  = 'Normalized'

## We use this to overwrite some of the settings above on a local basis if needed
readSettings('Accounting', AccountingConfig)
readSettings('Database', DatabaseConfig)

#try:
#    from mysql import *
#except ImportError:
#    if AccountingConfig.accounting == 'database':
#        error("The MySQLdb module which is required to do accounting is missing")
#        raise
#try:
#    from radius import *
#except ImportError:
#    if AccountingConfig.accounting == 'radius':
#        error("The pyrad module which is required to do accounting is missing")
#        raise

ModulesPath = os.path.realpath((os.path.dirname(__file__) or '.'))
ProgramName = os.path.splitext(os.path.basename(sys.argv[0] or 'python'))[0]

class StopRecord(dict):
    def __init__(self, session):
        info = {
            'id':           session._id,
            'fromTag':      session.tags[0],
            'toTag':        session.tags[1],
            'callDuration': session.callDuration,
            'callerBytes':  session.bytes[0],
            'calledBytes':  session.bytes[1],
            'codecs':       ';'.join(session.codecs),
            'userAgent':    ' + '.join(session.userAgents),
            'streamType':   '/'.join(unique_sorted(session.streamTypes)),
            'mediaInfo':    ''
        }
        if session.didTimeout:
            if session.onHold:
                info['mediaInfo'] = 'timeout on hold'
            else:
                info['mediaInfo'] = 'timeout'
        dict.__init__(self, info)

class StopRecordSerializer(SafeSerializer):
    _allowed = {'accounting': ['StopRecord']}
    __codecs = ('zlib', 'base64', 'string_escape')
    def __encode(self, stream):
        s = stream
        for codec in self.__codecs:
            s = s.encode(codec)
        return s
    def __decode(self, stream):
        s = stream
        for codec in self.__codecs[::-1]:
            s = s.decode(codec)
        return s
    def dump(self, obj):
        return self.__encode(SafeSerializer.dump(self, obj))
    def load(self, stream):
        return SafeSerializer.load(self, self.__decode(stream))

class MySQLAccounting(MySQL):
    __query = """
        UPDATE
            %(table)s
        SET
            %(durationField)s = %(durationField)s + IF(%(stopInfoField)s IS NULL, %%(callDuration)s, 0),
            %(stopTimeField)s = DATE_ADD(%(startTimeField)s, INTERVAL %(durationField)s SECOND),
            %(callerBytesField)s = %(callerBytesField)s + %%(callerBytes)s,
            %(calledBytesField)s = %(calledBytesField)s + %%(calledBytes)s,
            %(userAgentField)s = %%(userAgent)s,
            %(codecTypesField)s = %%(codecs)s,
            %(streamTypeField)s = %%(streamType)s,
            %(mediaInfoField)s = %%(mediaInfo)s,
            %(normalizedField)s = '0'
        WHERE
            %(sessionIdField)s = %%(id)s AND %(fromTagField)s = %%(fromTag)s AND %(toTagField)s = %%(toTag)s
        """ % DatabaseConfig.__dict__

    def makequery(self, record):
        try:
            return self.__query % self.literal(dict(record))
        except KeyError, which:
            raise RuntimeError("can't build MySQL log query: missing key: %s" % which)

    def log(self, record):
        try:
            changed = self.query(self.makequery(record))
            if changed == 0:
                record['toTag'] = 'n/a'
                self.query(self.makequery(record))
        except RuntimeError, msg:
            error("accounting failed: runtime error: " + msg)
        except MySQLError, why:
            error("accounting failed. MySQL error: %s" % (why,))
            error("failed query: %s" % self._query)


class RadiusAccounting(RadiusClient):
    def make_attributes(self, record):
        try:
            return {'Acct-Status-Type':  'Update',
                    'User-Name':         'mediaproxy@default',
                    'Acct-Session-Id':   record['id'],
                    'Acct-Session-Time': record['callDuration'],
                    'Acct-Input-Octets':  record['callerBytes'],
                    'Acct-Output-Octets': record['calledBytes'],
                    'Sip-From-Tag': record['fromTag'],
                    'Sip-To-Tag':   record['toTag'],
                    'Sip-User-Agents':  record['userAgent'],
                    'Sip-Applications': record['streamType'],
                    'Media-Codecs': record['codecs'],
                    'Media-Info':   record['mediaInfo']}
        except KeyError, which:
            raise RuntimeError("can't build radius accounting packet: missing key: %s" % which)

    def log(self, record):
        try:
            attributes = self.make_attributes(record)
            self.do_accounting(**attributes)
        except RuntimeError, msg:
            error("accounting failed: runtime error: " + str(msg))
        except RadiusAttributeError, msg:
            error("accounting failed: radius dictionary attribute error: " + str(msg))
        except RadiusValueError, msg:
            error("accounting failed: radius dictionary value error: " + str(msg))
        except TimeoutError, msg:
            error("accounting failed: radius server did timeout: " + str(msg))
        except PacketError, msg:
            error("accounting failed. radius packet error: %s" % str(msg))


class QueuedItemProcessingThread(threading.Thread):
    '''
    Provide a queue that will process the received items in a separate thread.
    This is useful where item processing is expected to take some time,
    while the main thread cannot block, thus a separate thread is used.

    The class provides the ability to save the unprocessed items on exit in a
    specified file and restore them when the program starts again, putting the
    items back in the queue for processing.

    To use this class, subclass it and overwrite the process method. The
    subclass' __init__ method should also call the parent initialization with
    the appropriate name and file arguments that are suitable for the subclass
    '''

    def __init__(self, name='QueuedItemProcessing', file=None):
        threading.Thread.__init__(self, name=name)
        self.__queue = Queue.Queue()
        self.__exit = threading.Event()
        if file is not None:
            self.__file = '%s/%s' % (ModulesPath, file)
        else:
            self.__file = None
        self.__restore()

    def __atexit(self):
        self.__exit.set()      ## signal thread that we want to exit
        self.__queue.put(None) ## add a dummy entry to wake up the thread
        self.join()            ## wait until the thread finishes processing the queue

    def __restore(self):
        try:
            f = open(self.__file, 'r')
        except:
            pass
        else:
            try:
                items = cPickle.load(f)
            except:
                items = []
            f.close()
            try:    os.unlink(self.__file)
            except: pass
            for item in items:
                self.__queue.put(item)

    def __save(self, entries):
        if not entries:
            return
        try:
            f = open(self.__file, 'w')
        except:
            pass
        else:
            failedDump = False
            try:
                try:
                    cPickle.dump(entries, f)
                except:
                    failedDump = True
            finally:
                f.close()
            if failedDump:
                try:    os.unlink(self.__file)
                except: pass

    def start(self):
        atexit.register(self.__atexit)
        threading.Thread.start(self)

    def run(self):
        while True:
            item = self.__queue.get()
            if self.__exit.isSet():
                break
            if item is None:
                continue
            ## process item here
            try:
                self.process(item)
            except:
                ## don't let this thread die no matter what.
                ## print the traceback of any uncaptured error
                info = sys.exc_info()
                try:
                    traceback.print_exception(*info)
                finally:
                    del info
        ## We are about to exit. Save remaining unprocessed items
        if item is not None:
            rest = [item]
        else:
            rest = []
        while True:
            try:
                item = self.__queue.get_nowait()
            except Queue.Empty:
                break
            else:
                if item is not None:
                    rest.append(item)
        ## save entries to disk
        self.__save(rest)

    def put(self, item):
        self.__queue.put(item)

    def process(self, item):
        '''Function that does the actual item processing. Must be overwritten in subclasses'''
        raise RuntimeError, "class %s doesn't define a `process' method" % self.__class__.__name__


class DatabaseAccountingThread(QueuedItemProcessingThread):
    def __init__(self):
        ## Safeguard for both dispatcher and mediaproxy using this class at the same time
        filename = '.%s-accounting.dat' % ProgramName
        QueuedItemProcessingThread.__init__(self, name='DatabaseAccounting', file=filename)
        self.__db = MySQLAccounting(shared=False)

    def process(self, record):
        self.__db.log(record)

    def log(self, session):
        if isinstance(session, StopRecord):
            if not session['toTag']:
                return ## there is no to-tag => session was not started. nothing to do.
            self.put(session)
        else:
            if not session.tags[1]:
                return ## there is no to-tag => session was not started. nothing to do.
            self.put(StopRecord(session))

class RadiusAccountingThread(QueuedItemProcessingThread):
    def __init__(self):
        ## Safeguard for both dispatcher and mediaproxy using this class at the same time
        filename = '.%s-accounting.dat' % ProgramName
        QueuedItemProcessingThread.__init__(self, name='RadiusAccounting', file=filename)
        self.__radius = RadiusAccounting()

    def process(self, record):
        self.__radius.log(record)

    def log(self, session):
        if isinstance(session, StopRecord):
            if not session['toTag']:
                return ## there is no to-tag => session was not started. nothing to do.
            self.put(session)
        else:
            if not session.tags[1]:
                return ## there is no to-tag => session was not started. nothing to do.
            self.put(StopRecord(session))

try:
    accounting
except NameError:
    if AccountingConfig.accounting == 'database':
        accounting = DatabaseAccountingThread()
        accounting.start()
    elif AccountingConfig.accounting == 'radius':
        accounting = RadiusAccountingThread()
        accounting.start()
    else:
        accounting = Null()

