"""Enhanced Beancount data structures for Conservancy

The classes in this module are interface-compatible with Beancount's core data

structures, and provide additional business logic that we want to use

throughout Conservancy tools.

"""

# Copyright © 2020  Brett Smith

#

# This program is free software: you can redistribute it and/or modify

# it under the terms of the GNU Affero General Public License as published by

# the Free Software Foundation, either version 3 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU Affero General Public License for more details.

#

# You should have received a copy of the GNU Affero General Public License

# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import collections

import datetime

import decimal

import functools

from beancount.core import account as bc_account

from beancount.core import amount as bc_amount

from beancount.core import convert as bc_convert

from typing import (

    cast,

    Callable,

    Iterable,

    Iterator,

    MutableMapping,

    Optional,

    Sequence,

    Union,

)

from .beancount_types import (

    Directive,

    MetaKey,

    MetaValue,

    Posting as BasePosting,

    Transaction,

)

DecimalCompat = Union[decimal.Decimal, int]

LINK_METADATA = frozenset([

    'approval',

    'check',

    'contract',

    'invoice',

    'purchase-order',

    'receipt',

    'rt-id',

    'statement',

])

class Account(str):

    """Account name string

    This is a string that names an account, like Assets:Bank:Checking

    or Income:Donations. This class provides additional methods for common

    account name parsing and queries.

    """

    __slots__ = ()

    SEP = bc_account.sep

    def is_cash_equivalent(self) -> bool:

        return (

            self.is_under('Assets:') is not None

            and self.is_under('Assets:Prepaid', 'Assets:Receivable') is None

        )

    def is_checking(self) -> bool:

        return self.is_cash_equivalent() and ':Check' in self

    def is_credit_card(self) -> bool:

        return self.is_under('Liabilities:CreditCard') is not None

    def is_opening_equity(self) -> bool:

        return self.is_under('Equity:Funds', 'Equity:OpeningBalance') is not None

    def is_under(self, *acct_seq: str) -> Optional[str]:

        """Return a match if this account is "under" a part of the hierarchy

        Pass in any number of account name strings as arguments. If this

        account is under one of those strings in the account hierarchy, the

        first matching string will be returned. Otherwise, None is returned.

        You can use the return value of this method as a boolean if you don't

        care which account string is matched.

        An account is considered to be under itself:

          Account('Expenses:Tax').is_under('Expenses:Tax') # returns 'Expenses:Tax'

    # within easy reach.

    @property

    def date(self) -> datetime.date:

        return self.txn.date

class Posting(BasePosting):

    """Enhanced Posting objects

    This class is a subclass of Beancount's native Posting class where

    specific fields are replaced with enhanced versions:

    * The `account` field is an Account object

    * The `units` field is our Amount object (which simply declares that the

      number is always a Decimal—see that docstring for details)

    * The `meta` field is a PostingMeta object

    """

    __slots__ = ()

    account: Account

    units: Amount

    # mypy correctly complains that our MutableMapping is not compatible

    # with Beancount's meta type declaration of Optional[Dict]. IMO

    # Beancount's type declaration is a smidge too specific: I think its type

    # declaration should also use MutableMapping, because it would be very

    # unusual for code to specifically require a Dict over that.

    # If it did, this declaration would pass without issue.

    meta: PostingMeta  # type:ignore[assignment]

    @classmethod

    def from_beancount(cls,

                       txn: Transaction,

                       index: int,

                       post: Optional[BasePosting]=None,

    ) -> 'Posting':

        if post is None:

            post = txn.postings[index]

        return cls(

            Account(post.account),

            *post[1:5],

            # see rationale above about Posting.meta

            PostingMeta(txn, index, post), # type:ignore[arg-type]

        )

    @classmethod

    def from_txn(cls, txn: Transaction) -> Iterable['Posting']:

        """Yield an enhanced Posting object for every posting in the transaction"""

        for index, post in enumerate(txn.postings):

            yield cls.from_beancount(txn, index, post)

    @classmethod

    def from_entries(cls, entries: Iterable[Directive]) -> Iterable['Posting']:

        """Yield an enhanced Posting object for every posting in these entries"""

        for entry in entries:

            # Because Beancount's own Transaction class isn't type-checkable,

            # we can't statically check this. Might as well rely on duck

            # typing while we're at it: just try to yield postings from

            # everything, and ignore entries that lack a postings attribute.

            try:

                yield from cls.from_txn(entry)  # type:ignore[arg-type]

            except AttributeError:

                pass

def balance_of(txn: Transaction,

               *preds: Callable[[Account], Optional[bool]],

) -> Amount:

    """Return the balance of specified postings in a transaction.

    Given a transaction and a series of account predicates, balance_of

    returns the balance of the amounts of all postings with accounts that

    match any of the predicates.

    balance_of uses the "weight" of each posting, so the return value will

    use the currency of the postings' cost when available.

    """

    match_posts = [post for post in Posting.from_txn(txn)

                   if any(pred(post.account) for pred in preds)]

    number = decimal.Decimal(0)

    if not match_posts:

        currency = ''

    else:

        weights: Sequence[Amount] = [

            bc_convert.get_weight(post) for post in match_posts

        ]

        number = sum((wt.number for wt in weights), number)

        currency = weights[0].currency

    return Amount(number, currency)

def is_opening_balance_txn(txn: Transaction) -> bool:

    opening_equity = balance_of(txn, Account.is_opening_equity)

    if not opening_equity.currency:

"""Test data.is_opening_balance_txn function"""

# Copyright © 2020  Brett Smith

#

# This program is free software: you can redistribute it and/or modify

# it under the terms of the GNU Affero General Public License as published by

# the Free Software Foundation, either version 3 of the License, or

# (at your option) any later version.

#

# This program is distributed in the hope that it will be useful,

# but WITHOUT ANY WARRANTY; without even the implied warranty of

# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

# GNU Affero General Public License for more details.

#

# You should have received a copy of the GNU Affero General Public License

# along with this program.  If not, see <https://www.gnu.org/licenses/>.

from decimal import Decimal

import pytest

from . import testutil

from conservancy_beancount import data

def test_typical_opening():

    assert data.is_opening_balance_txn(txn)

def test_multiacct_opening():

    txn = testutil.Transaction(postings=[

        ('Assets:Receivable:Accounts', 100),

        (next(testutil.OPENING_EQUITY_ACCOUNTS), -100),

        ('Liabilities:Payable:Accounts', -150),

        (next(testutil.OPENING_EQUITY_ACCOUNTS), 150),

    ])

    assert data.is_opening_balance_txn(txn)

def test_opening_with_fx():

    equity_post = txn.postings[-1]

    txn.postings[-1] = equity_post._replace(

        units=testutil.Amount(equity_post.units.number * Decimal('.9'), 'EUR'),

        cost=testutil.Cost('1.11111'),

    )

    assert data.is_opening_balance_txn(txn)

@pytest.mark.parametrize('acct1,acct2,number', [

    ('Assets:Receivable:Accounts', 'Income:Donations', 100),

    ('Expenses:Other', 'Liabilities:Payable:Accounts', 200),

    ('Expenses:Other', 'Equity:Retained:Costs', 300),

    # Release from restriction

    ('Equity:Funds:Unrestricted', 'Equity:Funds:Restricted', 400),

])

def test_not_opening_balance(acct1, acct2, number):

    txn = testutil.Transaction(postings=[

        (acct1, number),

        (acct2, -number),

    ])

    assert not data.is_opening_balance_txn(txn)

        ('Expenses:General', 25),

    ])

    errors = list(hook.run(txn))

    assert 1 <= len(errors) <= 2

    assert all(error.message == "transaction has invalid entity: {}".format(src_value)

               for error in hook.run(txn))

@pytest.mark.parametrize('src_value', VALID_VALUES)

def test_valid_values_on_payee(hook, src_value):

    txn = testutil.Transaction(payee=src_value, postings=[

        ('Assets:Cash', -25),

        ('Expenses:General', 25),

    ])

    assert not any(hook.run(txn))

@pytest.mark.parametrize('src_value', INVALID_VALUES)

def test_invalid_values_on_payee(hook, src_value):

    txn = testutil.Transaction(payee=src_value, postings=[

        ('Assets:Cash', -25),

        ('Expenses:General', 25),

    ])

    errors = list(hook.run(txn))

    assert 1 <= len(errors) <= 2

    assert all(error.message == "transaction has invalid entity: {}".format(src_value)

               for error in hook.run(txn))

@pytest.mark.parametrize('payee,src_value', testutil.combine_values(

    INVALID_VALUES,

    VALID_VALUES,

))

def test_invalid_payee_but_valid_metadata(hook, payee, src_value):

    txn = testutil.Transaction(**{'payee': payee, TEST_KEY: src_value}, postings=[

        ('Assets:Cash', -25),

        ('Expenses:Other', 25),

    ])

    assert not any(hook.run(txn))

@pytest.mark.parametrize('account,required', [

    ('Assets:Bank:Checking', False),

    ('Assets:Cash', False),

    ('Assets:Receivable:Accounts', True),

    ('Assets:Receivable:Loans', True),

    ('Equity:OpeningBalances', False),

    ('Expenses:General', True),

    ('Income:Donations', True),

    ('Liabilities:CreditCard', False),

    ('Liabilities:Payable:Accounts', True),

    ('Liabilities:Payable:Vacation', True),

    ('Liabilities:UnearnedIncome:Donations', False),

])

def test_which_accounts_required_on(hook, account, required):

    txn = testutil.Transaction(postings=[

        ('Assets:Checking', -25),

        (account, 25),

    ])

    errors = list(hook.run(txn))

    if not required:

        assert not errors

    else:

        assert errors

        assert any(error.message == "{} missing entity".format(account)

                   for error in errors)

def test_not_required_on_opening(hook):

    assert not list(hook.run(txn))

        (acct1, 25, {TEST_KEY: value}),

    ])

    expected = {

        MISSING_MSG(acct1),

        WRONG_TYPE_MSG(acct1, type(value).__name__),

    }

    actual = {error.message for error in hook.run(txn)}

    assert actual == expected

@pytest.mark.parametrize('acct1,acct2,value', testutil.combine_values(

    REQUIRED_ACCOUNTS,

    NON_REQUIRED_ACCOUNTS,

    testutil.LINK_METADATA_STRINGS,

))

def test_valid_values_on_transaction(hook, acct1, acct2, value):

    txn = testutil.Transaction(**{TEST_KEY: value}, postings=[

        (acct2, -25),

        (acct1, 25),

    ])

    assert not list(hook.run(txn))

@pytest.mark.parametrize('acct1,acct2,value', testutil.combine_values(

    REQUIRED_ACCOUNTS,

    NON_REQUIRED_ACCOUNTS,

    testutil.NON_LINK_METADATA_STRINGS,

))

def test_invalid_values_on_transaction(hook, acct1, acct2, value):

    txn = testutil.Transaction(**{TEST_KEY: value}, postings=[

        (acct2, -25),

        (acct1, 25),

    ])

    actual = {error.message for error in hook.run(txn)}

    assert actual == {"{} missing {}".format(acct1, TEST_KEY)}

@pytest.mark.parametrize('acct1,acct2,value', testutil.combine_values(

    REQUIRED_ACCOUNTS,

    NON_REQUIRED_ACCOUNTS,

    testutil.NON_STRING_METADATA_VALUES,

))

def test_bad_type_values_on_transaction(hook, acct1, acct2, value):

    txn = testutil.Transaction(**{TEST_KEY: value}, postings=[

        (acct2, -25),

        (acct1, 25),

    ])

    expected = {

        MISSING_MSG(acct1),

        WRONG_TYPE_MSG(acct1, type(value).__name__),

    }

    actual = {error.message for error in hook.run(txn)}

    assert actual == expected

@pytest.mark.parametrize('acct1,acct2', testutil.combine_values(

    REQUIRED_ACCOUNTS,

    NON_REQUIRED_ACCOUNTS,

))

def test_missing_invoice(hook, acct1, acct2):

    txn = testutil.Transaction(postings=[

        (acct2, -25),

        (acct1, 25),

    ])

    actual = {error.message for error in hook.run(txn)}

    assert actual == {"{} missing {}".format(acct1, TEST_KEY)}

def test_not_required_on_opening(hook):

    assert not list(hook.run(txn))

    ('Liabilities:CreditCard', False),

    ('Liabilities:Payable:Accounts', True),

    # We do want a "project" for Lia:Pay:Vacation but it has a default value

    ('Liabilities:Payable:Vacation', False),

    ('Liabilities:UnearnedIncome:Donations', True),

])

def test_which_accounts_required_on(hook, account, required):

    txn = testutil.Transaction(postings=[

        ('Assets:Checking', 25),

        (account, 25),

    ])

    errors = list(hook.run(txn))

    assert required == any(errors)

@pytest.mark.parametrize('account', [

    'Expenses:Payroll:Salary',

    'Expenses:Payroll:Tax',

    'Liabilities:Payable:Vacation',

])

def test_default_values(hook, account):

    txn = testutil.Transaction(postings=[

        ('Assets:Checking', -25),

        (account, 25),

    ])

    errors = list(hook.run(txn))

    assert not errors

    testutil.check_post_meta(txn, None, {TEST_KEY: DEFAULT_VALUE})

@pytest.mark.parametrize('date,required', [

    (testutil.EXTREME_FUTURE_DATE, False),

    (testutil.FUTURE_DATE, True),

    (testutil.FY_START_DATE, True),

    (testutil.FY_MID_DATE, True),

    (testutil.PAST_DATE, None),

])

def test_default_value_set_in_date_range(hook, date, required):

    txn = testutil.Transaction(date=date, postings=[

        ('Expenses:Payroll:Benefits', 25),

        ('Liabilities:Payable:Vacation', -25),

    ])

    errors = list(hook.run(txn))

    assert not errors

    expect_meta = {TEST_KEY: DEFAULT_VALUE} if required else None

    testutil.check_post_meta(txn, expect_meta, expect_meta)

@pytest.mark.parametrize('repo_path', [

    None,

    '..',

])

def test_missing_project_data(repo_path):

    config = testutil.TestConfig(repo_path=repo_path)

    with pytest.raises(errormod.ConfigurationError):

        meta_project.MetaProject(config)

@pytest.mark.parametrize('repo_path_s,data_path_s', [

    ('repository', 'Projects/project-list.yml'),

    ('..', 'LICENSE.txt'),

])

def test_invalid_project_data(repo_path_s, data_path_s):

    config = testutil.TestConfig(repo_path=repo_path_s)

    with pytest.raises(errormod.ConfigurationError):

        meta_project.MetaProject(config, Path(data_path_s))

def test_not_required_on_opening(hook):

    assert not list(hook.run(txn))

def test_always_required_on_restricted_funds(hook):

    acct = 'Equity:Funds:Restricted'

    actual = {error.message for error in hook.run(txn)}

    assert actual == {f'{acct} missing project'}

PAST_DATE = datetime.date(2000, 1, 1)

TESTS_DIR = Path(__file__).parent

def check_post_meta(txn, *expected_meta, default=None):

    assert len(txn.postings) == len(expected_meta)

    for post, expected in zip(txn.postings, expected_meta):

        if not expected:

            assert not post.meta

        else:

            actual = None if post.meta is None else {

                key: post.meta.get(key, default) for key in expected

            }

            assert actual == expected

def combine_values(*value_seqs):

    stop = 0

    for seq in value_seqs:

        try:

            stop = max(stop, len(seq))

        except TypeError:

            pass

    return itertools.islice(

        zip(*(itertools.cycle(seq) for seq in value_seqs)),

        stop,

    )

def date_seq(date=FY_MID_DATE, step=1):

    while True:

        yield date

        date += datetime.timedelta(days=step)

def parse_date(s, fmt='%Y-%m-%d'):

    return datetime.datetime.strptime(s, fmt).date()

def test_path(s):

    if s is None:

        return s

    s = Path(s)

    if not s.is_absolute():

        s = TESTS_DIR / s

    return s

def Amount(number, currency='USD'):

    return bc_amount.Amount(Decimal(number), currency)

def Cost(number, currency='USD', date=FY_MID_DATE, label=None):

    return bc_data.Cost(Decimal(number), currency, date, label)

def Posting(account, number,

            currency='USD', cost=None, price=None, flag=None,

            type_=bc_data.Posting, **meta):

    if cost is not None:

        cost = Cost(*cost)

    if not meta:

        meta = None

    return type_(

        account,

        Amount(number, currency),

        cost,

        price,

        flag,

        meta,

    )

LINK_METADATA_STRINGS = {

    'Invoices/304321.pdf',

    'rt:123/456',

    'rt://ticket/234',

}

NON_LINK_METADATA_STRINGS = {

    '',

    ' ',

    '     ',

}

NON_STRING_METADATA_VALUES = [

    Decimal(5),

    FY_MID_DATE,

    Amount(50),

    Amount(500, None),

]

OPENING_EQUITY_ACCOUNTS = itertools.cycle([

    'Equity:Funds:Unrestricted',

    'Equity:Funds:Restricted',

    'Equity:OpeningBalance',

])

def balance_map(source=None, **kwargs):

    # The source and/or kwargs should map currency name strings to

    # things you can pass to Decimal (a decimal string, an int, etc.)

    # This returns a dict that maps currency name strings to Amount instances.

    retval = {}

    if source is not None:

        retval.update((currency, Amount(number, currency))

                      for currency, number in source)

    if kwargs:

        retval.update(balance_map(kwargs.items()))

    return retval

class TestConfig:

    def __init__(self, *,

                 payment_threshold=0,

                 repo_path=None,

                 rt_client=None,

    ):

        self._payment_threshold = Decimal(payment_threshold)

        self.repo_path = test_path(repo_path)

        self._rt_client = rt_client

        if rt_client is None:

            self._rt_wrapper = None

        else:

            self._rt_wrapper = rtutil.RT(rt_client)

    def payment_threshold(self):

        return self._payment_threshold

    def repository_path(self):

        return self.repo_path

    def rt_client(self):

        return self._rt_client

    def rt_wrapper(self):

        return self._rt_wrapper

class _TicketBuilder:

    MESSAGE_ATTACHMENTS = [

        ('(Unnamed)', 'multipart/alternative', '0b'),

        ('(Unnamed)', 'text/plain', '1.2k'),

        ('(Unnamed)', 'text/html', '1.4k'),

    ]

    MISC_ATTACHMENTS = [

        ('Forwarded Message.eml', 'message/rfc822', '3.1k'),

        ('photo.jpg', 'image/jpeg', '65.2k'),

        ('ConservancyInvoice-301.pdf', 'application/pdf', '326k'),

        ('Company_invoice-2020030405_as-sent.pdf', 'application/pdf', '50k'),

        ('statement.txt', 'text/plain', '652b'),

        ('screenshot.png', 'image/png', '1.9m'),

    ]

    def __init__(self):

        self.id_seq = itertools.count(1)

        self.misc_attchs = itertools.cycle(self.MISC_ATTACHMENTS)

    def new_attch(self, attch):

        return (str(next(self.id_seq)), *attch)

    def new_msg_with_attachments(self, attachments_count=1):

        for attch in self.MESSAGE_ATTACHMENTS:

            yield self.new_attch(attch)

        for _ in range(attachments_count):

            yield self.new_attch(next(self.misc_attchs))

    def new_messages(self, messages_count, attachments_count=None):

        for n in range(messages_count):

            if attachments_count is None:

                att_count = messages_count - n

            else:

                att_count = attachments_count

            yield from self.new_msg_with_attachments(att_count)