#!/usr/bin/env python
"""
Read an OFX file and output its entries as Ledger syntax.

This program also reads a mapping from account id to account name, from the
ledger file itself, which should include lines of this form::

   @var ofx 000067632326            Assets:Current:RBC:Checking
   @var ofx 000023245336            Assets:Current:RBC:US-Checking
   @var ofx 000018783275            Assets:Current:RBC:Savings
   @var ofx 3233762676464639        Liabilities:Credit-Card:RBC-VISA

(Note that this is not part of the normal Ledger format, which is why the
directives are in comments.)
"""

# stdlib imports
import os, re, logging
from datetime import datetime, date, timedelta
from decimal import Decimal
from collections import defaultdict
from os.path import *
from xml.sax.saxutils import unescape

# beancount imports
from beancount import cmdline
from beancount.ledger import Transaction, Posting, Account
from beancount.wallet import Wallet
from beancount.ledger import read_ofx_accounts_map

# other imports
from BeautifulSoup import BeautifulStoneSoup



class OFXAccount(object):
    "An OFX account object."

    account = None
    currency = None
    transactions = None
    start, end = None, None
    bal_amount = None
    bal_time = None

    def __init__(self):
        self.transactions = []


class OFXTransaction(object):
    "An OFX transaction object."

    trntype = None
    dtposted = None
    trnamt = None
    fitid = None
    name = None
    memo = None
    checknum = None

    # Wallet with the amount of the transaction, and the other side (if available).
    amount = None
    amount_other = None

    def initialize(self, currency):
        if self.dtposted:
            self.dtposted = ofx_parse_time(self.dtposted).date()
        if self.memo is not None and re.match('[0-9.]+ [A-Z]{3} @ [0-9.]+', self.memo):
            self.amount_other, self.memo = self.memo, None
        self.description = u' -- '.join(filter(None, [self.name, self.memo]))
        self.amount = Wallet(currency, Decimal(self.trnamt))

    def __str__(self):
        date_s = self.dtposted.strftime('%Y-%m-%d')
        if self.checknum:
            return '%s ! (%s) %s' % (date_s, self.checknum,
                                     self.description)
        else:
            return '%s ! %s' % (date_s, self.description)


def process_ofx(fn, accounts_map):
    "Read an OFX file and process it. Return a list of account objects."
    all = []
    soup = BeautifulStoneSoup(open(fn))
    for st in soup.findAll(['stmttrnrs', 'ccstmttrnrs']):
        ofxacc = OFXAccount()

        accid, ofxacc.currency = find_account(st)
        try:
            ofxacc.account = accounts_map[accid]
        except KeyError:
            raise KeyError("Missing account name in map for %s" % accid)

        # Find transactions list.
        tranlist = st.find('banktranlist')

        # Find start and end times.
        ofxacc.start = ofx_parse_time(tranlist.dtstart.contents[0])
        ofxacc.end = ofx_parse_time(tranlist.dtend.contents[0])

        # Process transactions themselves.
        for trn in tranlist.findAll('stmttrn'):
            t = OFXTransaction()
            ofxacc.transactions.append(t)
            for a in 'trntype dtposted trnamt fitid name memo checknum'.split():
                n = trn.find(a)
                if n:
                    setattr(t, a, unescape(n.contents[0].strip()))
            t.initialize(ofxacc.currency)

        # Find expected balance.
        ofxacc.bal_amount = Decimal(st.ledgerbal.balamt.contents[0].strip())
        ofxacc.bal_time = ofx_parse_time(st.ledgerbal.dtasof.contents[0].strip())

        all.append(ofxacc)

    return all

def ofx_parse_time(s):
    "Parse an OFX time string and return a datetime object.."
    if len(s) < 14:
        return datetime.strptime(s[:8], '%Y%m%d')
    else:
        return datetime.strptime(s[:14], '%Y%m%d%H%M%S')

def find_account(st):
    """
    Get and return the account information and currency that was found in the
    text file.
    """
    s = st.find(['stmtrs', 'ccstmtrs'])
    assert s is not None
    acct = s.find(['bankacctfrom', 'ccacctfrom'])
    # Note: RBC offers a malformed XML.
    acctid = acct.acctid.contents[0].strip()
    currency = s.curdef.contents[0].strip()
    return acctid, currency




def main():
    import optparse
    parser = optparse.OptionParser(__doc__.strip())
    opts, ledger, args = cmdline.main(parser, 1)

    if not args:
        parser.error("You must provide a list of OFX files to convert.")

    accounts_map = read_ofx_accounts_map(ledger)

    # Create a map of postings by date.
    bydate = defaultdict(list)
    for post in ledger.postings:
        bydate[post.actual_date].append(post)

    oneday = timedelta(days=1)
    datesearch_begin_offset = timedelta(days=-3)
    datesearch_end_offset = timedelta(days=3)


    for arg in args:
        print '\n'*3
        print ';;; Import of %s' % arg

        for ofxacc in process_ofx(arg, accounts_map):
            account = ofxacc.account

            # Calculate boundaries within which we skip the transactions as
            # already entered (because they are located within the largest valid
            # check interval). We only skip when we have a non-empty range of
            # checks.
            check_min = None
            if (account.check_min is not None and
                account.check_max is not None and
                account.check_min != account.check_max):

                check_min = account.check_min
                check_max = account.check_max

            print '\n'*3
            hdfmt = '; %s import  %s  %s'
            print hdfmt % ('Start', account.fullname, ofxacc.start)
            print
            for t in ofxacc.transactions:

                # Skip entries between the largest valid check interval.
                if (check_min is not None and check_min < t.dtposted < check_max):
                    logging.warning("Skipping %s" % t)
                    continue

                print str(t)
                print '  %-60s %s' % (account.fullname, t.amount)
                if t.amount_other is not None:
                    print '  %-60s %s %s' % ('?', t.amount_other, ofxacc.currency)

                # Look for potential matches in the ledger file.
                date_ = t.dtposted + datesearch_begin_offset
                date_end = t.dtposted + datesearch_end_offset
                while date_ <= date_end:
                    for post in bydate.get(date_, []):
                        if t.amount == post.amount:
                            print ';; Potential match at  %s:%s' % (abspath(post.filename), post.lineno)
                    date_ += oneday
                print

            print hdfmt % ('End', account.fullname, ofxacc.end)
            if re.search('\\bCredit-Card:HSBC\\b', account.fullname):
                ofxacc.bal_amount = -ofxacc.bal_amount
            print '@check %s %s     %s %s' % (
                ofxacc.bal_time.date(), account.fullname, ofxacc.bal_amount, ofxacc.currency)
            print

    print '\n'


## FIXME: we also need to try to match up transactions between those that are
## being imported... this means that we should do it in a two step process.


if __name__ == '__main__':
    main()

