#!/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)