__copyright__ = "Copyright (C) 2013-2017, 2019-2020, 2024  Martin Blais"
__license__ = "GNU GPLv2"

import datetime
import pickle
import unittest
from datetime import date

from beancount.core import data
from beancount.core.amount import A

META = data.new_metadata("beancount/core/testing.beancount", 12345)
FLAG = "*"


class TestData(unittest.TestCase):
    def create_empty_transaction(self):
        return data.Transaction(
            META,
            date(2014, 1, 15),
            FLAG,
            None,
            "Some example narration",
            data.EMPTY_SET,
            data.EMPTY_SET,
            [],
        )

    def test_create_simple_posting(self):
        entry = self.create_empty_transaction()
        posting = data.create_simple_posting(entry, "Assets:Bank:Checking", "123.45", "USD")
        self.assertTrue(isinstance(posting, data.Posting))
        self.assertTrue(entry.postings)
        self.assertEqual(entry.postings[0], posting)

    def test_create_simple_posting_with_cost(self):
        entry = self.create_empty_transaction()
        posting = data.create_simple_posting_with_cost(
            entry, "Assets:Bank:Checking", "100", "MSFT", "123.45", "USD"
        )
        self.assertTrue(isinstance(posting, data.Posting))
        self.assertTrue(entry.postings)
        self.assertEqual(entry.postings[0], posting)

    def test_sanity_check_types(self):
        entry = self.create_empty_transaction()
        data.sanity_check_types(entry)
        with self.assertRaises(AssertionError):
            data.sanity_check_types("a string")
        with self.assertRaises(AssertionError):
            data.sanity_check_types(META)
        with self.assertRaises(AssertionError):
            data.sanity_check_types(datetime.date.today())
        with self.assertRaises(AssertionError):
            data.sanity_check_types(entry._replace(flag=1))
        with self.assertRaises(AssertionError):
            data.sanity_check_types(entry._replace(payee=1))
        with self.assertRaises(AssertionError):
            data.sanity_check_types(entry._replace(narration=1))
        with self.assertRaises(AssertionError):
            data.sanity_check_types(entry._replace(tags={}))
        with self.assertRaises(AssertionError):
            data.sanity_check_types(entry._replace(links={}))
        with self.assertRaises(AssertionError):
            data.sanity_check_types(entry._replace(postings=None))

    def test_posting_has_conversion(self):
        entry = self.create_empty_transaction()
        posting = data.create_simple_posting(entry, "Assets:Bank:Checking", "123.45", "USD")
        self.assertFalse(data.posting_has_conversion(posting))
        posting = posting._replace(price=A("153.02 CAD"))
        self.assertTrue(data.posting_has_conversion(posting))

    def test_transaction_has_conversion(self):
        entry = self.create_empty_transaction()
        posting = data.create_simple_posting(entry, "Assets:Bank:Checking", "123.45", "USD")
        posting = posting._replace(price=A("153.02 CAD"))
        entry.postings[0] = posting
        self.assertTrue(data.transaction_has_conversion(entry))

    def test_get_entry(self):
        entry = self.create_empty_transaction()
        posting = data.create_simple_posting(entry, "Assets:Bank:Checking", "123.45", "USD")
        self.assertEqual(entry, data.get_entry(entry))
        self.assertEqual(entry, data.get_entry(data.TxnPosting(entry, posting)))

    def create_sort_data(self):
        account = "Assets:Bank:Checking"
        date1 = date(2014, 1, 15)
        date2 = date(2014, 1, 18)
        date3 = date(2014, 1, 20)
        entries = [
            data.Transaction(
                data.new_metadata(".", 1100), date3, FLAG, None, "Next day", None, None, []
            ),
            data.Close(data.new_metadata(".", 1000), date2, account),
            data.Balance(
                data.new_metadata(".", 1001), date2, account, A('200.00 USD"'), None, None
            ),
            data.Open(data.new_metadata(".", 1002), date2, account, "USD", None),
            data.Transaction(
                data.new_metadata(".", 1009),
                date2,
                FLAG,
                None,
                "Transaction 2",
                None,
                None,
                [],
            ),
            data.Transaction(
                data.new_metadata(".", 1008),
                date2,
                FLAG,
                None,
                "Transaction 1",
                None,
                None,
                [],
            ),
            data.Transaction(
                data.new_metadata(".", 900),
                date1,
                FLAG,
                None,
                "Previous day",
                None,
                None,
                [],
            ),
        ]

        for entry in entries:
            if isinstance(entry, data.Transaction):
                data.create_simple_posting(entry, "Assets:Bank:Checking", "123.45", "USD")

        return entries

    def check_sorted(self, entries):
        self.assertEqual(
            [
                data.Transaction,
                data.Open,
                data.Balance,
                data.Transaction,
                data.Transaction,
                data.Close,
                data.Transaction,
            ],
            list(map(type, entries)),
        )

        self.assertEqual(
            [900, 1002, 1001, 1008, 1009, 1000, 1100],
            [entry.meta["lineno"] for entry in entries],
        )

    def test_entry_sortkey(self):
        entries = self.create_sort_data()
        sorted_entries = sorted(entries, key=data.entry_sortkey)
        self.check_sorted(sorted_entries)

    def test_sort(self):
        entries = self.create_sort_data()
        sorted_entries = data.sorted(entries)
        self.check_sorted(sorted_entries)

    def test_posting_sortkey(self):
        entries = self.create_sort_data()
        txn_postings = [
            (
                data.TxnPosting(entry, entry.postings[0])
                if isinstance(entry, data.Transaction)
                else entry
            )
            for entry in entries
        ]
        sorted_txn_postings = sorted(txn_postings, key=data.posting_sortkey)

        self.assertEqual(
            [
                data.TxnPosting,
                data.Open,
                data.Balance,
                data.TxnPosting,
                data.TxnPosting,
                data.Close,
                data.TxnPosting,
            ],
            list(map(type, sorted_txn_postings)),
        )

        self.assertEqual(
            [900, 1002, 1001, 1008, 1009, 1000, 1100],
            [entry.meta["lineno"] for entry in map(data.get_entry, sorted_txn_postings)],
        )

    def test_filter_txns(self):
        entries = self.create_sort_data()
        txns = list(data.filter_txns(entries))
        self.assertEqual(4, len(txns))
        self.assertTrue(all(isinstance(txn, data.Transaction) for txn in txns))

    def test_has_entry_account_component(self):
        entry = data.Transaction(
            data.new_metadata(".", 0),
            datetime.date.today(),
            FLAG,
            None,
            "Something",
            None,
            None,
            [],
        )
        data.create_simple_posting(entry, "Liabilities:US:CreditCard", "-50", "USD")
        data.create_simple_posting(entry, "Expenses:Food:Restaurant", "50", "USD")

        has_component = data.has_entry_account_component
        self.assertTrue(has_component(entry, "US"))
        self.assertFalse(has_component(entry, "CA"))

        self.assertTrue(has_component(entry, "CreditCard"))
        self.assertTrue(has_component(entry, "Liabilities"))
        self.assertFalse(has_component(entry, "Assets"))

        self.assertTrue(has_component(entry, "Restaurant"))
        self.assertTrue(has_component(entry, "Food"))
        self.assertTrue(has_component(entry, "Expenses"))
        self.assertFalse(has_component(entry, "Equity"))

    def test_find_closest(self):
        entry1 = data.Transaction(
            data.new_metadata("/tmp/apples.beancount", 200),
            datetime.date(2014, 9, 14),
            "*",
            None,
            "",
            None,
            None,
            [],
        )

        # Insert a decoy from another file (should fail).
        entry2 = data.Transaction(
            data.new_metadata("/tmp/bananas.beancount", 100),
            datetime.date(2014, 9, 20),
            "*",
            None,
            "",
            None,
            None,
            [],
        )

        entry3 = data.Transaction(
            data.new_metadata("/tmp/apples.beancount", 105),
            datetime.date(2014, 10, 1),
            "*",
            None,
            "",
            None,
            None,
            [],
        )

        entries = [entry1, entry2, entry3]

        # Try an exact match.
        self.assertTrue(data.find_closest(entries, "/tmp/apples.beancount", 105) is entry3)

        # Try a bit after.
        self.assertTrue(data.find_closest(entries, "/tmp/apples.beancount", 107) is entry3)

        # Try way closer to the next one.
        self.assertTrue(data.find_closest(entries, "/tmp/apples.beancount", 199) is entry3)

        # Get the next one.
        self.assertTrue(data.find_closest(entries, "/tmp/apples.beancount", 201) is entry1)

        # Get from the other file.
        self.assertTrue(data.find_closest(entries, "/tmp/bananas.beancount", 300) is entry2)

        # Get none.
        self.assertTrue(data.find_closest(entries, "/tmp/apples.beancount", 99) is None)

    def test_remove_account_postings(self):
        meta = data.new_metadata(".", 0)
        date = datetime.date.today()
        entry1 = data.Open(meta, date, "Liabilities:US:CreditCard", None, None)
        entry2 = data.Open(meta, date, "Equity:Rounding", None, None)

        entry3 = data.Transaction(meta, date, FLAG, None, "Something", None, None, [])
        data.create_simple_posting(entry3, "Liabilities:US:CreditCard", "-50", "USD")
        data.create_simple_posting(entry3, "Equity:Rounding", "0.00123", "USD")
        data.create_simple_posting(entry3, "Expenses:Food:Restaurant", "50", "USD")

        entry4 = data.Price(meta, date, "HOOL", A("23 USD"))

        in_entries = [entry1, entry2, entry3, entry4]

        out_entries = data.remove_account_postings("Equity:Rounding", in_entries)
        self.assertEqual(4, len(out_entries))
        self.assertEqual(
            ["Liabilities:US:CreditCard", "Expenses:Food:Restaurant"],
            [posting.account for posting in out_entries[2].postings],
        )

    def test_iter_entry_dates(self):
        prototype = data.Transaction(
            data.new_metadata("misc", 200), None, "*", None, "", None, None, []
        )
        dates = [
            datetime.date(2016, 1, 10),
            datetime.date(2016, 1, 11),
            datetime.date(2016, 1, 14),
            datetime.date(2016, 1, 15),
            datetime.date(2016, 1, 16),
            datetime.date(2016, 1, 20),
            datetime.date(2016, 1, 22),
        ]

        entries = [prototype._replace(date=date) for date in dates]

        # Same date, present.
        self.assertEqual(
            [],
            [
                entry.date
                for entry in data.iter_entry_dates(
                    entries, datetime.date(2016, 1, 14), datetime.date(2016, 1, 14)
                )
            ],
        )
        # Same date, absent.
        self.assertEqual(
            [],
            [
                entry.date
                for entry in data.iter_entry_dates(
                    entries, datetime.date(2016, 1, 12), datetime.date(2016, 1, 12)
                )
            ],
        )
        # Both dates exist.
        self.assertEqual(
            [datetime.date(2016, 1, 11), datetime.date(2016, 1, 14)],
            [
                entry.date
                for entry in data.iter_entry_dates(
                    entries, datetime.date(2016, 1, 11), datetime.date(2016, 1, 15)
                )
            ],
        )
        # First doesn't exist.
        self.assertEqual(
            [datetime.date(2016, 1, 14)],
            [
                entry.date
                for entry in data.iter_entry_dates(
                    entries, datetime.date(2016, 1, 12), datetime.date(2016, 1, 15)
                )
            ],
        )
        # Second doesn't exist.
        self.assertEqual(
            [datetime.date(2016, 1, 11)],
            [
                entry.date
                for entry in data.iter_entry_dates(
                    entries, datetime.date(2016, 1, 11), datetime.date(2016, 1, 13)
                )
            ],
        )
        # Neither exist.
        self.assertEqual(
            [
                datetime.date(2016, 1, 14),
                datetime.date(2016, 1, 15),
                datetime.date(2016, 1, 16),
            ],
            [
                entry.date
                for entry in data.iter_entry_dates(
                    entries, datetime.date(2016, 1, 13), datetime.date(2016, 1, 17)
                )
            ],
        )
        # Before.
        self.assertEqual(
            [datetime.date(2016, 1, 10)],
            [
                entry.date
                for entry in data.iter_entry_dates(
                    entries, datetime.date(2016, 1, 5), datetime.date(2016, 1, 11)
                )
            ],
        )
        # After.
        self.assertEqual(
            [datetime.date(2016, 1, 22)],
            [
                entry.date
                for entry in data.iter_entry_dates(
                    entries, datetime.date(2016, 1, 21), datetime.date(2016, 1, 30)
                )
            ],
        )
        # Around.
        self.assertEqual(
            dates,
            [
                entry.date
                for entry in data.iter_entry_dates(
                    entries, datetime.date(2016, 1, 2), datetime.date(2016, 1, 30)
                )
            ],
        )


class TestPickle(unittest.TestCase):
    def test_data_tuples_support_pickle(self):
        txn1 = data.Transaction(
            META,
            date(2014, 1, 15),
            FLAG,
            None,
            "Some example narration",
            data.EMPTY_SET,
            data.EMPTY_SET,
            [],
        )
        pickled_str = pickle.dumps(txn1)
        txn2 = pickle.loads(pickled_str)
        self.assertEqual(txn1, txn2)


if __name__ == "__main__":
    unittest.main()
