import click
import tomli as toml
import petl
import petlutils

from collections import defaultdict
from copy import copy
from itertools import groupby
from os import path
from orders import list_designated_accounts

from beancount import loader
from beancount.core import amount, data, inventory


def units(values):
    assert len(values) == 1
    position = values[0].get_only_position()
    if position is None:
        return amount.from_string('0.00 EUR')
    return position.units


def cumulative_balances_by_year(entries, accounts, unpack=False):
    balances = defaultdict(inventory.Inventory)
    for year, txns in groupby(data.filter_txns(entries), key=lambda entry: entry.date.year):
        for entry in txns:
            for posting in entry.postings:
                if posting.account in accounts:
                    balances[posting.account].add_position(posting)
        if unpack:
            for account, balance in balances.items():
                yield year, account, copy(balance)
        else:
            yield year, {account: copy(balance) for account, balance in balances.items()}


@click.command()
def main():
    with open(path.expanduser('~/.config/orders'), 'rb') as f:
        conf = toml.load(f)
    ledger = path.normpath(path.expanduser(conf['ledger']))

    entries, errors, options = loader.load_file(ledger)
    assert not errors

    accounts = set(list_designated_accounts(entries))
    balances = cumulative_balances_by_year(entries, accounts, unpack=True)

    table = petl.pushheader(balances, ['year', 'account', 'balance']) \
                .pivot('account', 'year', 'balance', units, missing='')

    # Fix data columns width.
    table = table.convert(table.fieldnames()[1:], lambda x: '{:>14s}'.format(str(x)))

    print(table)


if __name__ == '__main__':
    main()