#!/usr/bin/env python3

import re

from beancount.core import data
from beancount.core import inventory
from beancount.core import account
from beancount.core import amount
from beancount.core import account_types
from beancount.parser import parser
from beancount.core.number import D
from beancount.parser import options
import itertools as it
from collections import namedtuple

LinkAccount = namedtuple("LinkAccount", ["account", "tags"])

__plugins__ = ('link_account_plugin',)

def add_tags(txn: data.Transaction, *tags: str) -> data.Transaction:
    return txn._replace(tags=txn.tags | set(tags))

def link_posting(p: data.Posting, link_accounts: dict[str, LinkAccount]) -> list[data.Posting]:
    if p.account in link_accounts:
        return [
            p._replace(account=link_accounts[p.account].account),
            p._replace(units=-p.units)
        ]
    return None
    

def link_txn(txn: data.Transaction, link_accounts: dict[str, LinkAccount]) -> list[data.Transaction]:
    new_postings = [p for p in txn.postings if p.account in link_accounts]
    if len(new_postings) == 0:
        return []
    
    new_txns = []

    get_account = lambda x: x.account
    new_postings = sorted(new_postings, key=get_account)
    for k, g in it.groupby(new_postings, key=get_account):
        new_postings = [p for o in g for p in link_posting(o, link_accounts)]
        new_txns.append(add_tags(txn._replace(postings=new_postings), "link", *link_accounts[k].tags))
    return new_txns

def link_account_plugin(entries, options_map, config=None):
    errors = []

    link_accounts = {}
    for entry in entries:
        if isinstance(entry, data.Open):
            if "link-account" in entry.meta:
                tags = entry.meta["link-tags"].split(" ") if "link-tags" in entry.meta else []
                link_accounts[entry.account] = LinkAccount(account=entry.meta["link-account"], tags=tags)
    
    new_entries = []
    for entry in entries:
        if isinstance(entry, data.Transaction):
            new_entries.extend(link_txn(entry, link_accounts))
        new_entries.append(entry)

    return new_entries, errors