#!/usr/bin/env python3
from beancount import loader
from beanquery import query
from beancount.parser import printer
import argparse
from tabulate import tabulate
from decimal import Decimal
from beancount.core.amount import Amount, add, sub, mul
from datetime import date
from dateutil.relativedelta import relativedelta


class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


def draw_line():
    print('─' * 30)


def get_income_val(obj, key):
    if key in obj:
        amount = obj[key].get_only_position().units
        amount = mul(amount, Decimal(-1.0))
        return amount.to_string()
    else:
        return None


def get_expense_val(obj, key):
    if key in obj:
        amount = obj[key].get_only_position().units
        return amount.to_string()
    else:
        return None


def print_expenses_table(expenses):
    table = []
    for key, expense in expenses.items():
        parts = key.split(":", 1)
        table.append([parts[1], get_expense_val(expenses, key)])
    print(tabulate(table))


def get_total_inflows(income):
    sum = 0
    for account, balance in income.items():
        sum = balance if sum == 0 else sum + balance
    if sum != 0 and sum.get_only_position() != None:
        result = sum.get_only_position().units
        return Amount(Decimal(round(result.number, 2) * Decimal(-1.0)), result.currency)
    else:
        return Amount(Decimal(0), "EUR")


def get_total_outflows(expenses, total_investments):
    sum = 0
    for account, balance in expenses.items():
        sum = balance if sum == 0 else sum + balance
    if sum != 0 and sum.get_only_position() != None:
        result = sum.get_only_position().units
        return add(Amount(Decimal(round(result.number, 2)), result.currency), total_investments)
    else:
        return total_investments


def get_total_investments(investments):
    sum = Amount(Decimal(0), "EUR")
    for inv in investments:
        sum = inv.cost_position if sum == Amount(
            Decimal(0), "EUR") else add(sum, inv.cost_position)
    if sum != 0 and sum != None:
        return Amount(Decimal(round(sum.number, 2)), sum.currency)
    else:
        return Amount(Decimal(0), "EUR")


def print_report(start_date, period, income, expenses, investments):
    print(f"{bcolors.BOLD}Cash Flow Statement (period={
          period}, start_date={start_date}){bcolors.ENDC}")
    draw_line()
    print(f"{bcolors.BOLD}Inflows{bcolors.ENDC}")
    print(f"\t{bcolors.BOLD}Income{bcolors.ENDC}")
    print(tabulate([
        ["Salari", get_income_val(income, "Income:Work:Zurich:Salari")],
        ["Tickets Restaurant", get_income_val(
            income, "Income:Work:Zurich:TicketsRestaurant")],
        ["Targeta Transport", get_income_val(
            income, "Income:Work:Zurich:TargetaTransport")],
        ["Pla pensions Empleats Zurich", get_income_val(
            income, "Income:Work:Zurich:DZP")],
        ["Seguro Mèdic", get_income_val(
            income, "Income:Work:Zurich:SeguroMedic")],
        ["Gimnàs", get_income_val(income, "Income:Work:Zurich:Gimnas")]
    ]))
    print(f"\t{bcolors.BOLD}Income from Investment{bcolors.ENDC}")
    print(tabulate([
        ["Capital Gains", get_income_val(
            income, "Income:Invest:R4:CapitalGains")],
        ["Untaxable Capital Gains", get_income_val(
            income, "Income:Invest:R4:CapitalGains:Untaxable")],
        ["Dividends", get_income_val(income, "Income:Invest:R4:Dividends")],
        ["Rentabilitat Estalvis", get_income_val(
            income, "Income:Savings:Caixabank:RentabilitatEstalvis")]
    ]))
    print(f"\t{bcolors.BOLD}Other Inflows{bcolors.ENDC}")
    print(tabulate([
        ["Transferències", get_income_val(
            income, "Income:Other:Caixabank:Transferencia")],
        ["Bizum", get_income_val(income, "Income:Other:Caixabank:Bizum")],
        ["Devolucions", get_income_val(income, "Income:Other:Devolucions")]
    ]))
    print(tabulate([
        ["Total Inflows", f"{bcolors.BOLD}{get_total_inflows(income).to_string()}{
            bcolors.ENDC}"]
    ]))
    draw_line()
    print(f"{bcolors.BOLD}Outflows{bcolors.ENDC}")
    print_expenses_table(expenses)
    print(f"{bcolors.BOLD}Outflows from Investment{bcolors.ENDC}")
    total_investments = get_total_investments(investments)
    print(tabulate([
        ["Compra de fons i accions", total_investments.to_string()]
    ]))
    print(tabulate([
        ["Total Outflows", f"{bcolors.BOLD}{get_total_outflows(
            expenses, total_investments).to_string()}{bcolors.ENDC}"]
    ]))
    draw_line()
    net_cash_flow = sub(get_total_inflows(income),
                        get_total_outflows(expenses, total_investments))
    print(tabulate([
        ["NET CASH FLOW", f"{bcolors.BOLD}{bcolors.OKGREEN if net_cash_flow.number >= 0 else bcolors.FAIL}{
            net_cash_flow.to_string()}{bcolors.ENDC}"]
    ]))


def get_income(entries, options, period, start_date):
    period_delta = relativedelta(
        months=1) if period == "monthly" else relativedelta(years=1)
    end_date = date.fromisoformat(start_date) + period_delta
    income_query = f"SELECT account, convert(sum(position), \"EUR\") as sum_position FROM OPEN ON {
        start_date} CLOSE ON {end_date.isoformat()} WHERE account ~ \"Income\""
    rtypes, rrows = query.run_query(
        entries, options, income_query)
    income = {}
    for row in rrows:
        income[row[0]] = row[1]
    return income


def get_expenses(entries, options, period, start_date):
    period_delta = relativedelta(
        months=1) if period == "monthly" else relativedelta(years=1)
    end_date = date.fromisoformat(start_date) + period_delta
    expenses_query = f"SELECT account, convert(sum(position), \"EUR\") as sum_position FROM OPEN ON {
        start_date} CLOSE ON {end_date.isoformat()} WHERE account ~ \"Expenses\""
    rtypes, rrows = query.run_query(
        entries, options, expenses_query)
    expenses = {}
    for row in rrows:
        expenses[row[0]] = row[1]
    return expenses


def get_investments(entries, options, period, start_date):
    period_delta = relativedelta(
        months=1) if period == "monthly" else relativedelta(years=1)
    end_date = date.fromisoformat(start_date) + period_delta
    expenses_query = f"SELECT account, convert(cost(position), \"EUR\") as cost_position, currency, date WHERE account ~ \"Assets:Invest:R4:\" AND NOT currency ~ '^(EUR|USD)' AND date >= {
        start_date} AND date < {end_date.isoformat()}"
    rtypes, rrows = query.run_query(
        entries, options, expenses_query)
    return rrows


def main():
    parser = argparse.ArgumentParser(description='Generate cash flow report')
    parser.add_argument('start_date', metavar='start_date', type=str, nargs=1,
                        help='Start date (end date will be one month after if monthly report or one year after if yearly report)')
    parser.add_argument('-p', metavar='period', type=str, choices=["monthly", "yearly"], default="monthly", required=False,
                        help='Period (monthly or yearly)')

    args = parser.parse_args()
    start_date = args.start_date[0]
    period = args.p

    filename = "ledger/main.beancount"
    entries, errors, options = loader.load_file(filename)

    if errors:
        printer.print_errors(errors)

    income = get_income(entries, options, period, start_date)
    expenses = get_expenses(entries, options, period, start_date)
    investments = get_investments(entries, options, period, start_date)
    print_report(start_date, period, income, expenses, investments)


main()
