#!/usr/bin/env python3

import json
from decimal import Decimal
from datetime import datetime
from collections import OrderedDict, defaultdict
from operator import itemgetter

class Transaction:
    description = ""
    amount = Decimal(0)

    def __init__(self, **args):
        self.budget = defaultdict(Decimal)
        self.expenses = defaultdict(Decimal)
        self.date = datetime.utcnow().replace(microsecond = 0)
        balance = None
        for arg, val in args.items():
            if arg == "splits":
                self.add_splits(val)
            elif arg == "budget":
                for env, amount in val: 
                    self.add_budget(env, amount)
            elif arg == "expenses":
                for env, amount in val:
                    self.add_expense(env, amount)
            elif arg in ("balance", "envelop"):
                balance = val
            else:
                setattr(self, arg, val)
        if balance is not None:
            self.rebalance(balance)

    def add_budget(self, envelop, amount):
        assert(amount > 0)
        self.budget[envelop] += amount

    def add_expense(self, envelop, amount):
        assert(amount > 0)
        self.expenses[envelop] += amount

    def add_splits(self, *splits):
        for env, amount in splits:
            if amount < 0:
                self.add_expense(env, -amount)
            else:
                self.add_budget(env, amount)

    def rebalance(self, envelop):
        b = self.balance
        if b < 0:
            self.add_expense(envelop, -b)
        else:
            self.add_budget(envelop, b)
        assert(self.balance == 0)

    @property
    def balance(self):
        return self.amount \
            - sum(v for v in self.budget.values()) \
            + sum(v for v in self.expenses.values())

    @property
    def splits(self):
        for env, amount in self.budget.items():
            yield (env, amount)
        for env, amount in self.expenses.items():
            yield (env, -amount)

    def rename_envelop(self, oldname, newname):
        if oldname in self.budget:
            self.budget[newname] += self.budget[oldname]
            del self.budget[oldname]
        if oldname in self.expenses:
            self.expenses[newname] += self.expenses[oldname]
            del self.expenses[oldname]

    def as_dict(self):
        def serialize(values):
            for e, v in sorted(values.items(), key = itemgetter(0)):
                d = OrderedDict.fromkeys(["envelop", "value"])
                d.update(envelop = e, value = float(v))
                yield d
        d = OrderedDict.fromkeys(["description", "date", "amount", "budget", "expenses"])
        d.update(
            description = self.description,
            date = self.date.isoformat(),
            amount = float(self.amount),
            budget = list(serialize(self.budget)),
            expenses = list(serialize(self.expenses)))
        return d

    @classmethod
    def from_dict(cls, d):
        try:
            date = datetime.strptime("%Y-%m-%dT%H:%M:%S", d["date"])
        except:
            date = datetime.utcnow().replace(microsecond = 0)
        return Transaction(
            description = d.get("description", ""),
            date = date,
            amount = Decimal(d.get("amount", Decimal(0))),
            budget = ((v["envelop"], v["value"]) for v in d.get("budget", [])),
            expenses = ((v["envelop"], v["value"]) for v in d.get("expenses", [])))

class Envelop:
    name = ""
    description = ""
    budget = Decimal(0)
    expenses = Decimal(0)

    def __init__(self, name = "", description = ""):
        self.name = name
        self.description = description

    @property
    def balance(self):
        return self.budget - self.expenses

    def update(self, amount):
        if (amount < 0):
            self.expenses -= amount
        else:
            self.budget += amount

class Budget:
    name = "Budget"
    date = datetime.utcnow().replace(microsecond = 0)

    def __init__(self):
        self.metaenvelops = defaultdict(str)
        self.transactions = []

    def load(self, fp):
        d = json.load(fp, parse_float = Decimal)
        self.metaenvelops = dict((e["name"], e["description"]) for e in d.get("envelops",[]))
        self.transactions = []
        add = self.add_transaction
        for t in d.get("transactions",[]): add(Transaction.from_dict(t))

    def save(self, fp, **args):
        d = OrderedDict.fromkeys(["name", "date", "envelops", "transactions"])
        envelops = list()
        for name, desc in sorted(self.metaenvelops.items(), key = itemgetter(0)):
            if name != "":
                env = OrderedDict.fromkeys(["name", "description"])
                env.update(name = name, description = desc)
                envelops.append(env)
        d.update(name = self.name,
                 date = self.date.isoformat(),
                 envelops = envelops,
                 transactions = list(t.as_dict() for t in self.transactions))
        json.dump(d, fp, allow_nan=False, **args)

    def add_transaction(self, t):
        for e in (s[0] for s in t.splits):
            if e not in self.metaenvelops:
                self.metaenvelops[e] = ""
        self.transactions.append(t)
        return t

    def set_envelop_description(self, name, description):
        self.metaenvelops[name] = description
    
    def rename_envelop(self, oldname, newname):
        for t in self.transactions:
            t.rename_envelop(oldname, newname)
        if oldname in self.metaenvelops:
            self.metaenvelops.setdefault(newname, self.metaenvelops[oldname])
            del self.metaenvelops[oldname]

    @property
    def envelops(self):
        total = Envelop("Total")
        unassigned = Envelop("Unassigned", "Unassigned balance")
        envelops = defaultdict(Envelop)
        for t in self.transactions:
            for name, value in t.splits:
                assert(name)
                envelops[name].update(value)
            unassigned.update(t.balance)
        for name, env in envelops.items():
            env.name = name
            env.description = self.metaenvelops.get(name,"")
            total.budget += env.budget
            total.expenses += env.expenses
            yield env
        yield total
        yield unassigned

    def report(self):
        for env in self.envelops:
            yield "| {:<20} | {:<50} | {:> 10.2f} | {:> 10.2f} | {:> 10.2f} |".format(env.name, env.description, env.budget, env.expenses, env.balance)

    def read_transaction(self, reader = input):
        name = reader("Description: ")
        amount = Decimal(reader("Amount: "))
        t = Transaction(name = name, amount = amount)

        while t.balance != 0:
            envelop = reader("Balance: {:> 10.2f}\nEnvelop: ".format(t.balance))
            if not envelop:
                break
            amount = reader("Amount:" )
            if amount == "":
                t.rebalance(envelop)
            else:
                amount = Decimal(amount)
                t.add_expense(envelop, -amount) if amount < 0 else t.add_budget(envelop, amount)
        keep = reader("Do you want to keep the transaction [y/N]? ")
        if keep and keep[0].upper() == 'Y':
            return t
        return None
    

def main(filename, action):
    import os
    import sys
    b = Budget()
    if action == "init":
        with open(filename, "w") as fp:
            b.save(fp)
    elif action == "report":
        with open(filename, "r") as fp:
            b.load(fp)
        for l in b.report():
            print(l)
    elif action == "add":
        with open(filename, "r") as fp:
            b.load(fp)
        t = b.read_transaction()
        if t:
            b.add_transaction(t)
            os.rename(os.path.realpath(filename), os.path.realpath(filename) + ".bak")
            with open(filename, "w") as fp:
                b.save(fp)
    else:
        print("bad command line option")

if __name__ == '__main__':
    import sys
    import os
    filename = sys.argv[1]
    action = sys.argv[2]
    main(filename, action)