#!/usr/bin/env python3

from beancount.core import data
from beancount.core import account
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__ = ('joint_account_plugin',)

# DEBT_ACCOUNT = "Deute"
SHARE_ACCOUNT = "Part"

def debt_account_str(acc: str, debtor: str, accs: account_types.AccountTypes) -> str:
    return account.join(accs.liabilities, account.sans_root(acc), debtor)

def share_account_str(acc: str, debtor: str) -> str:
    # return ":".join((account, SHARE_ACCOUNT + "-" + debtor))
    return acc # reimbursment in same account

def debt_account(posting: data.Posting, debtor: str, joint_accounts: dict, accs: account_types.AccountTypes) -> str:
    if debtor in joint_accounts.get(posting.account, {}):
        return debt_account_str(posting.account, debtor, accs)
    if account_types.is_income_statement_account(posting.account, accs):
        return share_account_str(posting.account, debtor)
    return ""

def is_splitable(posting: data.Posting, accs: account_types.AccountTypes):
    if posting.meta.get("joint", "") == "nosplit":
        return False
    return account_types.is_income_statement_account(posting.account, accs)

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 add_tag(txn: data.Transaction, tag: str) -> data.Transaction:
    return txn._replace(tags=txn.tags | set([tag]))

def prepare_posting(posting, user, joint_accounts, accs, n):
    return divide_and_flip(posting._replace(account=debt_account(posting, user, joint_accounts, accs), flag=None, meta=None), n)

def split_among_users(postings, users, joint_accounts, accs):
    new_postings = []
    for user in users:
        for posting in postings:
            new_postings.append(prepare_posting(posting, user, joint_accounts, accs, 1 + len(users)))
    return new_postings


def debt_transaction(txn: data.Transaction, joint_accounts: dict, new_accounts: set, accs: account_types.AccountTypes) -> None | data.Transaction:
    # Only two cases are processed:
    # a) all accounts are joint with the same people
    # b) a single joint account with splitable/unsplitable accounts

    # get how many joint accounts there are
    user_sets = [(idx, set(joint_accounts[posting.account])) for idx, posting in enumerate(txn.postings) if posting.account in joint_accounts]

    new_postings = []
    
    # no joint account -> nothing to do
    if len(user_sets) == 0:
        return None
    
    # single joint account, assume all expenses/income are split equally between users of the joint account
    elif len(user_sets) == 1:
        splitable_postings = [posting for posting in txn.postings if is_splitable(posting, accs)]
        if len(splitable_postings) == 0:
            return None

        # as not every posting may be splittable, debt of joint account might need to be rebalanced
        new_balance = realization.compute_postings_balance(splitable_postings)
        joint_posting = txn.postings[user_sets[0][0]]
        joint_posting = joint_posting._replace(units=-new_balance.get_currency_units(joint_posting.units.currency))
        new_postings = splitable_postings + [joint_posting]
    
    # all postings are from joint accounts, check they share same users and transfer debt equally
    elif len(user_sets) == len(txn.postings):
        # check all joint accounts share same users
        total_set = set().union(*map(lambda s: s[1], user_sets))
        if any((total_set != s[1] for s in user_sets)):
            return None
        new_postings = txn.postings
        
    # just transfer debt equally
    new_postings = split_among_users(new_postings, user_sets[0][1], joint_accounts, accs)
    new_accounts |= set((p.account for p in new_postings))
    return add_tag(txn._replace(postings=new_postings), "retorn")

def joint_account_plugin(entries, options_map):
    errors = []
    new_entries = []
    new_accounts = set()

    accs = options.get_account_types(options_map)

    joint_accounts = {}
    for entry in entries:
        if isinstance(entry, data.Open):
            if "joint" in entry.meta:
                joint_accounts[entry.account] = entry.meta["joint"].strip().split(" ")
    
    for entry in entries:
        if isinstance(entry, data.Transaction):
            if any((acc in joint_accounts for acc in getters.get_accounts([entry]))):
                new_entry = debt_transaction(entry, joint_accounts, new_accounts, accs)
                if new_entry is not None:
                    new_entries.append(new_entry)
        
        new_entries.append(entry)
    
    # add new_accounts at earliest date possible
    earliest_date = getters.get_min_max_dates(new_entries)[0]
    new_accounts -= set(getters.get_accounts((e for e in entries if isinstance(e, data.Open))))

    for idx, acc in enumerate(new_accounts):
        entry = data.Open(data.new_metadata("<joint_account>", idx), earliest_date, acc, None, None)
        new_entries.append(entry)
    
    # close associated accounts
    for entry in entries:
        if isinstance(entry, data.Close):
            for user in joint_accounts.get(entry.account, []):
                e = entry._replace(account=debt_account_str(entry.account, user, accs))
                new_entries.append(e)

    return new_entries, []

# def change_to_total_account(txn: data.Transaction, joint_accounts) -> data.Transaction:
#     return txn._replace(postings=list(map(lambda p: change_to_total_account_posting(p, joint_accounts), txn.postings)))

# def change_to_total_account_posting(posting: data.Posting, joint_accounts) -> data.Posting:
#     return posting._replace(account=total_account_str(posting.account) if posting.account in joint_accounts else posting.account)

# def transform_accounts(entry, joint_accounts):
#     if hasattr(entry, "account"):
#         if entry.account in joint_accounts:
#             return entry._replace(account=total_account_str(entry.account))
#     elif isinstance(entry, data.Transaction):
#         return change_to_total_account(entry, joint_accounts)
#     return entry