#!/usr/bin/env python3

import re
import collections
from typing import Sequence

from beancount.core import data
from beancount.core import inventory
from beancount.core import account
from beancount.core import amount
from beancount.core import getters
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

__plugins__ = ('split_account_plugin',)

ConfigError = collections.namedtuple("ConfigError", "source message entry")
MultipleAccountError = collections.namedtuple("MultipleAccountError", "source message entry")
SplitAccount = collections.namedtuple("SplitAccount", ["account", "users", "tag"])

class MultipleAccountException(Exception):
    def __init__(self, split_tags):
        self.split_tags = split_tags

def split_posting(
        posting: data.Posting, 
        split: SplitAccount, 
        accs: account_types.AccountTypes, 
        main: None | str
        ) -> list[data.Posting]:
    weights = None
    if "split" in posting.meta:
        assert(set(posting.meta["split"].split(" ")) <= set(split.users))
        users = posting.meta["split"].split(" ")
        weights = [1 if user in users else 0 for user in split.users]
    elif account.leaf(posting.account) in split.users:
        user = account.leaf(posting.account)
        weights = [1 if u == user else 0 for u in split.users]
    else:
        weights = [1 for _ in range(len(split.users))]
    
    total_weight = sum(weights)
    assert(total_weight > 0)

    new_postings = [
        posting._replace(
            units=amount.mul(posting.units, -D(w) / total_weight),
            account=account.join(split.account, split.users[idx]),
        )
        for idx, w in enumerate(weights)
    ]
    if main in split.users:
        main_idx =  split.users.index(main)
        new_postings.extend([
            posting._replace(
                units=amount.mul(posting.units, -D(w) / total_weight),
            )
            for idx, w in enumerate(weights)
            if idx != main_idx
        ])
        new_postings.extend([
            posting._replace(
                units=amount.mul(posting.units, D(w) / total_weight),
                account=account.join("Passiu", "Split", split.users[idx]),
            )
            for idx, w in enumerate(weights)
            if idx != main_idx
        ])

    return new_postings

def merge_postings(postings: Sequence[data.Posting]) -> list[data.Posting]:
    merged = []
    keyfunc = lambda p: p.account
    for acc, g in it.groupby(sorted(postings, key=keyfunc), key=keyfunc):
        inv = inventory.Inventory()
        for p in g:
            inv.add_position(p)

        for pos in inv.get_positions():
            merged.append(data.Posting(
                account=acc,
                units=pos.units,
                cost=pos.cost,
                price=None,
                flag=None,
                meta=None
            ))
    return merged


def split_txn(
        txn: data.Transaction, 
        split_accounts: dict[str, SplitAccount], 
        split_tags: dict[str, str], 
        accs: account_types.AccountTypes, 
        main: None | str
        ) -> data.Transaction:
    involved_accs = set(p.account for p in txn.postings if p.account in split_accounts)

    tags = list(tag for tag in txn.tags if tag in split_tags)
    split_accs = list(split_tags[tag] for tag in tags)
    if len(split_accs) > 1:
        raise MultipleAccountException(tags)
    elif len(split_accs) == 1:
        involved_accs.add(split_accs[0])
    if len(involved_accs) == 0:
        return txn
    
    # filter postings that involve parent accounts, this way we can preserve all asset buckets
    new_postings = [p for p in txn.postings if p.account not in involved_accs]

    split_postings = []
    for acc in involved_accs:
        ps = [(p._replace(account="<invalid>") if p.account in involved_accs else p) for p in txn.postings if p.account != acc]
        split_postings.extend(n for p in ps for n in split_posting(p, split_accounts[acc], accs, main))

    # merge split_postings grouped by account
    new_postings.extend(merge_postings(split_postings))

    return txn._replace(postings=new_postings)

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

    split_accounts = {}
    split_tags = {}
    for entry in entries:
        if isinstance(entry, data.Open):
            if "split" in entry.meta:
                tag, *users = entry.meta["split"].split(" ")
                users = list(set(users)) # remove repeats
                if len(users) < 2:
                    errors.append(ConfigError(
                        entry.meta,
                        f"Split account need at least two users, only {len(users)} found: {users}",
                        entry
                    ))
                    continue

                split_accounts[entry.account] = SplitAccount(users=users, account=entry.account, tag=tag)
                split_tags[tag] = entry.account
                for user in users:
                    new_entries.append(
                        data.Open(
                            meta=data.new_metadata("<split-account>", 0),
                            date=entry.date,
                            account=account.join(entry.account, user),
                            currencies=entry.currencies,
                            booking=entry.booking
                        )
                    )
    old_accounts = getters.get_accounts(new_entries) | getters.get_accounts(entries)

    new_accounts = set()
    for entry in entries:
        if isinstance(entry, data.Transaction):
            try:
                entry = split_txn(entry, split_accounts, split_tags, accs, config)
                new_accounts |= set((p.account for p in entry.postings))
            except MultipleAccountException as e:
                errors.append(MultipleAccountError(
                    entry.meta,
                    f"Transaction has multiple split tags: {e.split_tags}",
                    entry
                ))
        new_entries.append(entry)

    oldest_date = getters.get_min_max_dates(new_entries)[0]
    for acc in new_accounts - old_accounts:
        new_entries.append(
            data.Open(
                meta=data.new_metadata("<split-account>", 1),
                date=oldest_date,
                account=acc,
                currencies=None,
                booking=None
            )
        )

    return new_entries, errors