"""core.py - Common data classes for reporting functionality"""

# 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

from decimal import Decimal

from .. import data

from typing import (

    overload,

    Dict,

    List,

    Sequence,

    Union,

)

class RelatedPostings(Sequence[data.Posting]):

    """Collect and query related postings

    This class provides common functionality for collecting related postings

    and running queries on them: iterating over them, tallying their balance,

    etc.

    This class doesn't know anything about how the postings are related. That's

    entirely up to the caller.

    A common pattern is to use this class with collections.defaultdict

    to organize postings based on some key::

        report = collections.defaultdict(RelatedPostings)

        for txn in transactions:

            for post in Posting.from_txn(txn):

                if should_report(post):

                    key = post_key(post)

                    report[key].add(post)

    """

    def __init__(self) -> None:

        self._postings: List[data.Posting] = []

    @overload

    def __getitem__(self, index: int) -> data.Posting: ...

    @overload

    def __getitem__(self, s: slice) -> Sequence[data.Posting]: ...

    def __getitem__(self,

                    index: Union[int, slice],

    ) -> Union[data.Posting, Sequence[data.Posting]]:

        if isinstance(index, slice):

            raise NotImplementedError("RelatedPostings[slice]")

        else:

            return self._postings[index]

    def __len__(self) -> int:

        return len(self._postings)

    def add(self, post: data.Posting) -> None:

        self._postings.append(post)

    def balance(self) -> Sequence[data.Amount]:

        currency_balance: Dict[str, Decimal] = collections.defaultdict(Decimal)

        for post in self:

            currency_balance[post.units.currency] += post.units.number

        return [data.Amount(number, key) for key, number in currency_balance.items()]

"""test_reports_related_postings - Unit tests for RelatedPostings"""

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

import itertools

import pytest

from . import testutil

from conservancy_beancount import data

from conservancy_beancount.reports import core

_day_counter = itertools.count(1)

def next_date():

    return testutil.FY_MID_DATE + datetime.timedelta(next(_day_counter))

def txn_pair(acct, src_acct, dst_acct, amount, date=None, txn_meta={}, post_meta={}):

    if date is None:

        date = next_date()

    src_txn = testutil.Transaction(date=date, **txn_meta, postings=[

        (acct, amount, post_meta.copy()),

        (src_acct, -amount),

    ])

    dst_date = date + datetime.timedelta(days=1)

    dst_txn = testutil.Transaction(date=dst_date, **txn_meta, postings=[

        (acct, -amount, post_meta.copy()),

        (dst_acct, amount),

    ])

    return (src_txn, dst_txn)

def donation(amount, currency='USD', date=None, other_acct='Assets:Cash', **meta):

    if date is None:

        date = next_date()

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

        ('Income:Donations', -amount, currency, meta),

        (other_acct, amount, currency),

    ])

def test_balance():

    related = core.RelatedPostings()

    related.add(data.Posting.from_beancount(donation(10), 0))

    assert related.balance() == [testutil.Amount(-10)]

    related.add(data.Posting.from_beancount(donation(15), 0))

    assert related.balance() == [testutil.Amount(-25)]

    related.add(data.Posting.from_beancount(donation(20), 0))

    assert related.balance() == [testutil.Amount(-45)]

def test_balance_zero():

    related = core.RelatedPostings()

    related.add(data.Posting.from_beancount(donation(10), 0))

    related.add(data.Posting.from_beancount(donation(-10), 0))

    assert related.balance() == [testutil.Amount(0)]

def test_balance_multiple_currencies():

    related = core.RelatedPostings()

    related.add(data.Posting.from_beancount(donation(10, 'GBP'), 0))

    related.add(data.Posting.from_beancount(donation(15, 'GBP'), 0))

    related.add(data.Posting.from_beancount(donation(20, 'EUR'), 0))

    related.add(data.Posting.from_beancount(donation(25, 'EUR'), 0))

    assert set(related.balance()) == {

        testutil.Amount(-25, 'GBP'),

        testutil.Amount(-45, 'EUR'),

    }

def test_balance_multiple_currencies_one_zero():

    related = core.RelatedPostings()

    related.add(data.Posting.from_beancount(donation(10, 'EUR'), 0))

    related.add(data.Posting.from_beancount(donation(15, 'USD'), 0))

    related.add(data.Posting.from_beancount(donation(-10, 'EUR'), 0))

    assert set(related.balance()) == {

        testutil.Amount(-15, 'USD'),

        testutil.Amount(0, 'EUR'),

    }