#!/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")
UnweightedError = collections.namedtuple("UnweightedError", "source message entry")
SplitAccount = collections.namedtuple("SplitAccount", ["account", "users", "tag"])

SPLIT_WEIGHT_CUSTOM = "split-weights"

def parse_split_weights(weight_str: str, users: list[str]) -> dict[str, "D"]:
    weights = {}
    for part in weight_str.split(" "):
        user, sep, weight = part.partition(":")
        if not sep:
            raise ValueError(f"Malformed weight entry: {part!r}")
        weights[user] = D(weight)
    if not set(weights) <= set(users):
        raise ValueError(f"Unknown users in weights: {set(weights) - set(users)}")
    return weights


def weights_for_tag_at_date(split_weights, tag, date, users: list[str]):
    history = split_weights.get(tag)
    if not history:
        return None
    applicable = None
    for entry_date, weights in history:
        if entry_date <= date:
            applicable = weights
        else:
            break
    if applicable is None:
        return None
    return [applicable.get(u, D("0")) for u in users]

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

class UnweightedException(Exception):
    def __init__(self):
        pass

def split_posting(
        posting: data.Posting,
        split: SplitAccount,
        accs: account_types.AccountTypes,
        main: None | str,
        split_weights,
        date,
        ) -> list[data.Posting]:

    weights = weights_for_tag_at_date(split_weights, split.tag, date, split.users)
    if weights is None:
        weights = [1 for _ in range(len(split.users))]

    if "split" in posting.meta:
        assert(set(posting.meta["split"].split(" ")) <= set(split.users))
        users = posting.meta["split"].split(" ")
        weights = [w if user in users else 0 for (w, user) in zip(weights, 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]
    
    total_weight = sum(weights)
    if total_weight == 0:
        raise UnweightedException()

    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,
        split_weights,
        ) -> 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, split_weights, txn.date))

    # 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)

    split_weights = collections.defaultdict(list)
    for entry in entries:
        if isinstance(entry, data.Custom) and entry.type == SPLIT_WEIGHT_CUSTOM:
            if len(entry.values) != 2:
                errors.append(ConfigError(
                    entry.meta,
                    f"Split weight custom entry needs exactly 2 arguments (tag, weights), got {len(entry.values)}",
                    entry
                ))
                continue
            tag = entry.values[0].value
            weight_str = entry.values[1].value
            if tag not in split_tags:
                errors.append(ConfigError(
                    entry.meta,
                    f"Split weight custom entry refers to unknown split tag: {tag}",
                    entry
                ))
                continue
            users = split_accounts[split_tags[tag]].users
            try:
                weights = parse_split_weights(weight_str, users)
            except ValueError as e:
                errors.append(ConfigError(entry.meta, f"Invalid split weights for tag {tag}: {e}", entry))
                continue
            split_weights[tag].append((entry.date, weights))

    for tag in split_weights:
        split_weights[tag].sort(key=lambda dw: dw[0])

    new_accounts = set()
    for entry in entries:
        if isinstance(entry, data.Transaction):
            try:
                entry = split_txn(entry, split_accounts, split_tags, accs, config, split_weights)
                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
                ))
            except UnweightedException as e:
                errors.append(UnweightedError(
                    entry.meta,
                    f"Split transaction had no payer as all involved users had no contribution",
                    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