#!/usr/bin/env python3

import re

from beancount.core import data
from beancount.core import getters
from beancount.core import amount
from beancount.core import account_types
from beancount.core import realization
from beancount.core.number import D
from beancount.parser import options

__plugins__ = ('split_account_plugin',)

def get_group(tag, _pattern=re.compile(r"split-(.*)")):
    m = _pattern.match(tag)
    if m:
        return m.group(1)
    return ""

def is_splitable(posting: data.Posting, accs: account_types.AccountTypes) -> bool:
    if posting.meta.get("split", "") == "no":
        return False
    return account_types.is_income_statement_account(posting.account, accs)

def add_tag(txn: data.Transaction, tag: str) -> data.Transaction:
    return txn._replace(tags=txn.tags | set([tag]))

def divide_and_flip(posting: data.Posting, n: int = 2) -> data.Posting:
    """Return a new posting with n-th the amount and opposite sign."""
    assert(n>0)
    old_amount = posting.units    
    new_amount = amount.div(old_amount, D(-n))
    return posting._replace(units=new_amount)

def split_txn(txn: data.Transaction, split_accounts: dict, accs: account_types.AccountTypes) -> None | data.Transaction:
    # involves_account = any((acc in split_accounts.items() for acc in getters.get_accounts([entry])))

    groups = [g for g in ( get_group(tag) for tag in txn.tags ) if g in split_accounts]
    if len(groups) != 1:
        return None
    group = groups[0]
    
    has_ignore_meta = txn.meta.get("split", "") == "no"
    if has_ignore_meta:
        return None
    
    # gather all splitable expenses/income, and get their total balance
    new_postings = [posting for posting in txn.postings if is_splitable(posting, accs)]
    if len(new_postings) == 0:
        return None
    
    new_balance = realization.compute_postings_balance(new_postings)
    for position in new_balance:
        new_postings.append(data.Posting(
            account=split_accounts[group],
            units=-position.units,
            cost=position.cost,
            price=None,
            flag=None,
            meta=None
        ))

    new_postings = [ divide_and_flip(p) for p in new_postings ]
    
    return add_tag(txn._replace(postings=new_postings), "retorn")

def split_account_plugin(entries, options_map):
    errors = []
    new_entries = []
    accs = options.get_account_types(options_map)

    split_accounts = {}
    for entry in entries:
        if isinstance(entry, data.Open):
            if "split" in entry.meta:
                split_accounts[entry.meta["split"]] = entry.account
    
    for entry in entries:
        if isinstance(entry, data.Transaction):
            new_txn = split_txn(entry, split_accounts, accs)
            if new_txn is not None:
                new_entries.append(new_txn)
        new_entries.append(entry)
            
    return new_entries, errors