from budgettracker import Budget, Transaction, Decimal

class Expect:
    def __init__(self, *cmds, verbose = False):
        self._cmds = cmds
        self._index = 0
        self.verbose = verbose
    def __call__(self, prompt):
        assert(self._index < len(self._cmds))
        idx = self._index
        self._index += 1
        if self.verbose:
            print("{:<20} => {}".format(prompt, self._cmds[idx]))
        return self._cmds[idx]

class Test:
    def run(self):
        self.test_load_save()
        self.test_report(True)
        self.test_envelops()
        self.test_read_transaction()

    def test_read_transaction(self, verbose = False):
        cmds = [ "Salary", "400",
                 "Taxes", "100",
                 "Taxes", "-100",
                 "Grocery", "200" ]
        b = self.create_budget()
        t = b.read_transaction(Expect(verbose = verbose, *(cmds + ["", "n"])))
        assert(t is None)
        t = b.read_transaction(Expect(verbose = verbose, *(cmds + ["", "y"])))
        assert(t.amount == 400)
        assert(t.budget["Taxes"] == 100)
        assert(t.budget["Grocery"] == 200)
        assert(len(t.budget) == 2)
        assert(t.expenses["Taxes"] == 100)
        assert(len(t.expenses) == 1)
        assert(t.balance == 200)
        t = b.read_transaction(Expect(verbose = verbose, *(cmds + ["Hobbies", "", "y"])))
        assert(t.amount == 400)
        assert(t.budget["Taxes"] == 100)
        assert(t.budget["Grocery"] == 200)
        assert(t.budget["Hobbies"] == 200)
        assert(len(t.budget) == 3)
        assert(t.expenses["Taxes"] == 100)
        assert(len(t.expenses) == 1)
        assert(t.balance == 0)

    def create_budget(self):
        b = Budget()
        salary = Decimal(500)
        taxes = salary * Decimal('0.20')
        # Salary: 500 - 100 in taxes = 400
        # Budget: 200$ in Grocery, 100$ for Taxes
        # Expenses: Taxes 100$
        # Unassigned: 200$
        t = Transaction(name="Salary", amount=salary-taxes,
                        budget = [("Grocery", Decimal(200)),
                                  ("Taxes", taxes)],
                        expenses = [("Taxes", taxes)])
        b.add_transaction(t)
        # Gift: 100$
        # Budget: 100$ in Hobbies
        t = Transaction(name = "Gift", amount = Decimal(100), envelop = "Hobbies")
        b.add_transaction(t)
        # Costco: -75$
        # Expenses: 75$ in Grocery
        t = Transaction(name="Costco", amount=Decimal(-75))
        t.rebalance("Grocery")
        b.add_transaction(t)
        # Cinema: -50
        # Expenses: 50$ in Hobbies
        t = Transaction(name="Cinema", amount=Decimal(-50), envelop="Hobbies")
        assert(t.balance == 0)
        b.add_transaction(t)
        # Missing money: -20$
        # Unassigned: -20$
        t = Transaction(name="Missing", amount=Decimal(-20))
        b.add_transaction(t)
        # Total:
        # Budget: 200$ Grocery, 100$ Taxes, 100$ Hobbies = 400$
        # Expenses: 100$ Taxes, 75$ Grocery, 50$ Hobbies = 225$
        # Unassigned: 200$ Salary, -20$ Missing = 180$ 

        b.set_envelop_description("Grocery", "stuff for stomach")
        b.rename_envelop("Grocery", "Food")
        return b

    def test_load_save(self, verbose = False):
        b = self.create_budget()

        from io import StringIO
        fp = StringIO()
        b.save(fp)
        saved = fp.getvalue()
        if verbose: print(saved)

        b2 = Budget()
        b2.load(StringIO(saved))
        fp = StringIO()
        b2.save(fp)
        saved2 = fp.getvalue()
        if verbose: print(saved2)
        assert(saved == saved2)

    def test_envelops(self, verbose = False):
        b = self.create_budget()
        envelops = list(b.envelops)
        assert(len(envelops) == 5)
        total = envelops[-2]
        unassigned = envelops[-1]

        assert(total.budget == 400)
        assert(total.expenses == 225)
        assert(total.balance == 175)
        assert(unassigned.budget == 200)
        assert(unassigned.expenses == 20)
        assert(unassigned.balance == 180)

    def test_report(self, verbose = False):
        b = self.create_budget()
        for l in b.report():
            if verbose:
                print(l)

if __name__ == '__main__':
    Test().run()