#!/usr/bin/env python
# -*- coding: utf-8 -*-
# vim: ft=python et softtabstop=4 cinoptions=4 shiftwidth=4 ts=4 ai

# Copyright(C) 2010-2011  Romain Bignon
#
# This file is part of weboob.
#
# weboob is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# weboob is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with weboob. If not, see <http://www.gnu.org/licenses/>.

import os
import sys
import locale
import time
import logging
from weboob.core import Weboob, CallErrors
from weboob.capabilities.bank import CapBank
from weboob.tools.browser import BrowserIncorrectPassword

class BoobankMuninPlugin(object):
    def __init__(self):
        if 'weboob_path' in os.environ:
            self.weboob = Weboob(os.environ['weboob_path'])
        else:
            self.weboob = Weboob()
        self.monitored_accounts = None
        if 'boobank_monitored' in os.environ:
            self.monitored_accounts = os.environ['boobank_monitored'].split(' ')
        self.cache_expire = long(os.environ.get('boobank_cache_expire', 3600))
        self.add_coming = int(os.environ.get('boobank_add_coming', 1))
        self.cumulate = int(os.environ.get('boobank_cumulate', 1))
        self.cache = None

    def display_help(self):
        print 'boobank-munin is a plugin for munin'
        print ''
        print 'Copyright(C) 2010-2011 Romain Bignon'
        print ''
        print 'To use it, create a symlink /etc/munin/plugins/boobank to this script'
        print 'and add this section in /etc/munin/plugin-conf.d/munin-node:'
        print ''
        print '[boobank]'
        print 'user romain'
        print 'group romain'
        print 'env.HOME /home/romain'
        print '# The weboob directory path.'
        print 'env.weboob_path /home/romain/.config/weboob/'
        print '# Monitored accounts. If this parameter is missing, all accounts'
        print '# will be displayed.'
        print 'env.boobank_monitored 0125XXXXXXXXXXXX@bnporc 0125XXXXXXXXXXXX@bnporc'
        print '# To prevent mass connections to bank websites, results are cached.'
        print '# You can set here the expiration delay (in seconds).'
        print 'env.boobank_cache_expire 7200'
        print '# If enabled, coming operations are added to the value of accounts\''
        print '# balance.'
        print 'env.boobank_add_coming 1'
        print '# Cumulate accounts values'
        print 'env.boobank_cumulate 1'
        print ''
        print 'When you change configuration, you can use this command to reset cache:'
        print '$ boobank-munin --reset'

    def clear_cache(self):
        for name in ('boobank-munin', 'boobank-munin-config'):
            try:
                os.unlink(self.cachepath(name))
            except IOError:
                pass

    def cachepath(self, name):
        tmpdir = os.path.join(self.weboob.workdir, "munin")
        if not os.path.isdir(tmpdir):
            os.makedirs(tmpdir)

        return os.path.join(tmpdir, name)

    def check_cache(self, name):
        return self.print_cache(name, check=True)

    def print_cache(self, name, check=False):
        try:
            f = open(self.cachepath(name), 'r')
        except IOError:
            return False

        try:
            last = int(f.readline().strip())
        except ValueError:
            return False

        if check and (last + self.cache_expire) < time.time():
            return False

        for line in f.xreadlines():
            sys.stdout.write(line)
        return True

    def new_cache(self, name):
        os.umask(0077)
        new_name = '%s.new' % name
        filename = self.cachepath(new_name)
        try:
            f = open(filename, 'w')
        except IOError, e:
            print >>sys.stderr, 'Unable to create the cache file %s: %s' % (filename, e)
            return

        self.cache = f
        self.cache.write('%d\n' % time.time())

    def flush_cache(self):
        old_name = self.cache.name
        new_name = self.cache.name[:-4]
        self.cache.close()
        os.rename(old_name, new_name)

    def write_output(self, line):
        sys.stdout.write('%s\n' % line)
        if self.cache:
            self.cache.write('%s\n' % line)

    def config(self):
        if self.check_cache('boobank-munin-config'):
            return

        self.new_cache('boobank-munin-config')
        self.weboob.load_backends(CapBank)
        self.write_output('graph_title Bank accounts')
        self.write_output('graph_vlabel balance')
        self.write_output('graph_category weboob')
        self.write_output('graph_args -l 0')
        try:
            accounts = []
            if self.monitored_accounts is not None:
                d = {}
                for backend, account in self.weboob.do('iter_accounts'):
                    if self.monitored(account):
                        d['%s@%s' % (account.id, account.backend)] = account

                for id in self.monitored_accounts:
                    try:
                        accounts.append(d[id])
                    except KeyError:
                        pass
            else:
                accounts = reversed([a for b, a in self.weboob.do('iter_accounts')])

            first = True
            for account in accounts:
                id = self.account2id(account)
                type = 'STACK'
                if first:
                    type = 'AREA'
                    first = False
                self.write_output('%s.label %s' % (id, account.label.encode('iso-8859-15')))
                if self.cumulate:
                    self.write_output('%s.draw %s' % (id, type))
        except CallErrors, errors:
            self.print_errors(errors)
            self.print_cache('boobank-munin-config')
        else:
            self.flush_cache()

    def monitored(self, account):
        return not self.monitored_accounts or ('%s@%s' % (account.id, account.backend)) in self.monitored_accounts

    def account2id(self, account):
        return '%s_%s' % (account.backend, account.id)

    def print_errors(self, errors):
        for backend, err, backtrace in errors:
            print >>sys.stderr, (u'%s(%s): %s' % (type(err).__name__, backend.name, err)).encode(sys.stdout.encoding or locale.getpreferredencoding(), 'replace')
            if isinstance(err, BrowserIncorrectPassword):
                self.weboob.backends_config.edit_backend(backend.name, backend.NAME, {'_enabled': False})

    def execute(self):
        if self.check_cache('boobank-munin'):
            return

        self.new_cache('boobank-munin')
        self.weboob.load_backends(CapBank)
        try:
            for backend, account in self.weboob.do('iter_accounts'):
                if self.monitored(account):
                    balance = account.balance
                    if account.coming and self.add_coming:
                        balance += account.coming
                    self.write_output('%s.value %d' % (self.account2id(account), balance))
        except CallErrors, errors:
            self.print_errors(errors)
            self.print_cache('boobank-munin')
        else:
            self.flush_cache()

    def run(self):
        cmd = (len(sys.argv) > 1 and sys.argv[1]) or "execute"
        if cmd == 'execute':
            self.execute()
        elif cmd == 'config':
            self.config()
        elif cmd == 'autoconf':
            print 'no'
            sys.exit(1)
        elif cmd == 'suggest':
            sys.exit(1)
        elif cmd == 'help' or cmd == '-h' or cmd == '--help':
            self.display_help()
        elif cmd == 'reload' or cmd == '--reload' or \
             cmd == 'reset' or cmd == '--reset':
            self.clear_cache()

        if self.cache:
            self.cache.close()

        sys.exit(0)

if __name__ == '__main__':
    logging.basicConfig()
    BoobankMuninPlugin().run()
