budgeteer/postgres/ynab-import.go

package postgres

import (
	"context"
	"encoding/csv"
	"fmt"
	"io"
	"strings"
	"time"
	"unicode/utf8"

	"github.com/google/uuid"
)

type YNABImport struct {
	accounts       []Account
	payees         []Payee
	categories     []GetCategoriesRow
	categoryGroups []CategoryGroup
	queries        *Queries
	budgetID       uuid.UUID
}

func NewYNABImport(context context.Context, queries *Queries, budgetID uuid.UUID) (*YNABImport, error) {
	accounts, err := queries.GetAccounts(context, budgetID)
	if err != nil {
		return nil, err
	}

	payees, err := queries.GetPayees(context, budgetID)
	if err != nil {
		return nil, err
	}

	categories, err := queries.GetCategories(context, budgetID)
	if err != nil {
		return nil, err
	}

	categoryGroups, err := queries.GetCategoryGroups(context, budgetID)
	if err != nil {
		return nil, err
	}

	return &YNABImport{
		accounts:       accounts,
		payees:         payees,
		categories:     categories,
		categoryGroups: categoryGroups,
		queries:        queries,
		budgetID:       budgetID,
	}, nil
}

// ImportAssignments expects a TSV-file as exported by YNAB in the following format:
// "Month"	"Category Group/Category"	"Category Group"	"Category"	"Budgeted"	"Activity"	"Available"
// "Apr 2019"	"Income: Next Month"	"Income"	"Next Month"	0,00€	0,00€	0,00€
//
// Activity and Available are not imported, since they are determined by the transactions and historic assignments.
func (ynab *YNABImport) ImportAssignments(context context.Context, r io.Reader) error {
	csv := csv.NewReader(r)
	csv.Comma = '\t'
	csv.LazyQuotes = true

	csvData, err := csv.ReadAll()
	if err != nil {
		return fmt.Errorf("read from tsv: %w", err)
	}

	count := 0
	for _, record := range csvData[1:] {
		dateString := record[0]
		date, err := time.Parse("Jan 2006", dateString)
		if err != nil {
			return fmt.Errorf("parse date %s: %w", dateString, err)
		}

		categoryGroup, categoryName := record[2], record[3] // also in 1 joined by :
		category, err := ynab.GetCategory(context, categoryGroup, categoryName)
		if err != nil {
			return fmt.Errorf("get category %s/%s: %w", categoryGroup, categoryName, err)
		}

		amountString := record[4]
		amount, err := GetAmount(amountString, "0,00€")
		if err != nil {
			return fmt.Errorf("parse amount %s: %w", amountString, err)
		}

		if amount.Int.Int64() == 0 {
			continue
		}

		assignment := CreateAssignmentParams{
			Date:       date,
			CategoryID: category.UUID,
			Amount:     amount,
		}
		_, err = ynab.queries.CreateAssignment(context, assignment)
		if err != nil {
			return fmt.Errorf("save assignment %v: %w", assignment, err)
		}

		count++
	}

	fmt.Printf("Imported %d assignments\n", count)

	return nil
}

type Transfer struct {
	CreateTransactionParams
	TransferToAccount *Account
	FromAccount       string
	ToAccount         string
}

// ImportTransactions expects a TSV-file as exported by YNAB in the following format:
// "Account"	"Flag"	"Date"	"Payee"	"Category Group/Category"	"Category Group"	"Category"	"Memo"	"Outflow"	"Inflow"	"Cleared"
// "Cash"	""	"11.12.2021"	"Transfer : Checking"	""	""	""	"Brought to bank"	500,00€	0,00€	"Cleared"
func (ynab *YNABImport) ImportTransactions(context context.Context, r io.Reader) error {
	csv := csv.NewReader(r)
	csv.Comma = '\t'
	csv.LazyQuotes = true

	csvData, err := csv.ReadAll()
	if err != nil {
		return fmt.Errorf("read from tsv: %w", err)
	}

	var openTransfers []Transfer

	count := 0
	for _, record := range csvData[1:] {
		transaction, err := ynab.GetTransaction(context, record)
		if err != nil {
			return err
		}

		payeeName := record[3]
		// Transaction is a transfer
		if strings.HasPrefix(payeeName, "Transfer : ") {
			err = ynab.ImportTransferTransaction(context, payeeName, transaction.CreateTransactionParams,
				&openTransfers, transaction.Account, transaction.Amount)
		} else {
			err = ynab.ImportRegularTransaction(context, payeeName, transaction.CreateTransactionParams)
		}
		if err != nil {
			return err
		}

		count++
	}

	for _, openTransfer := range openTransfers {
		fmt.Printf("Saving unmatched transfer from %s to %s on %s over %f as regular transaction\n",
			openTransfer.FromAccount, openTransfer.ToAccount, openTransfer.Date, openTransfer.Amount.GetFloat64())
		_, err = ynab.queries.CreateTransaction(context, openTransfer.CreateTransactionParams)
		if err != nil {
			return fmt.Errorf("save transaction %v: %w", openTransfer.CreateTransactionParams, err)
		}
	}

	fmt.Printf("Imported %d transactions\n", count)

	return nil
}

type NewTransaction struct {
	CreateTransactionParams
	Account *Account
}

func (ynab *YNABImport) GetTransaction(context context.Context, record []string) (NewTransaction, error) {
	accountName := record[0]
	account, err := ynab.GetAccount(context, accountName)
	if err != nil {
		return NewTransaction{}, fmt.Errorf("get account %s: %w", accountName, err)
	}

	// flag := record[1]

	dateString := record[2]
	date, err := time.Parse("02.01.2006", dateString)
	if err != nil {
		return NewTransaction{}, fmt.Errorf("parse date %s: %w", dateString, err)
	}

	categoryGroup, categoryName := record[5], record[6] // also in 4 joined by :
	category, err := ynab.GetCategory(context, categoryGroup, categoryName)
	if err != nil {
		return NewTransaction{}, fmt.Errorf("get category %s/%s: %w", categoryGroup, categoryName, err)
	}

	memo := record[7]

	outflow := record[8]
	inflow := record[9]
	amount, err := GetAmount(inflow, outflow)
	if err != nil {
		return NewTransaction{}, fmt.Errorf("parse amount from (%s/%s): %w", inflow, outflow, err)
	}

	statusEnum := TransactionStatusUncleared
	status := record[10]
	switch status {
	case "Cleared":
		statusEnum = TransactionStatusCleared
	case "Reconciled":
		statusEnum = TransactionStatusReconciled
	case "Uncleared":
	}

	return NewTransaction{
		CreateTransactionParams: CreateTransactionParams{
			Date:       date,
			Memo:       memo,
			AccountID:  account.ID,
			CategoryID: category,
			Amount:     amount,
			Status:     statusEnum,
		},
		Account: account,
	}, nil
}

func (ynab *YNABImport) ImportRegularTransaction(context context.Context, payeeName string,
	transaction CreateTransactionParams) error {
	payeeID, err := ynab.GetPayee(context, payeeName)
	if err != nil {
		return fmt.Errorf("get payee %s: %w", payeeName, err)
	}
	transaction.PayeeID = payeeID

	_, err = ynab.queries.CreateTransaction(context, transaction)
	if err != nil {
		return fmt.Errorf("save transaction %v: %w", transaction, err)
	}
	return nil
}

func (ynab *YNABImport) ImportTransferTransaction(context context.Context, payeeName string,
	transaction CreateTransactionParams, openTransfers *[]Transfer,
	account *Account, amount Numeric) error {
	transferToAccountName := payeeName[11:]
	transferToAccount, err := ynab.GetAccount(context, transferToAccountName)
	if err != nil {
		return fmt.Errorf("get transfer account %s: %w", transferToAccountName, err)
	}

	transfer := Transfer{
		transaction,
		transferToAccount,
		account.Name,
		transferToAccountName,
	}

	found := false
	for i, openTransfer := range *openTransfers {
		if openTransfer.TransferToAccount.ID != transfer.AccountID {
			continue
		}
		if openTransfer.AccountID != transfer.TransferToAccount.ID {
			continue
		}
		if openTransfer.Amount.GetFloat64() != -1*transfer.Amount.GetFloat64() {
			continue
		}

		fmt.Printf("Matched transfers from %s to %s over %f\n", account.Name, transferToAccount.Name, amount.GetFloat64())
		transfers := *openTransfers
		transfers[i] = transfers[len(transfers)-1]
		*openTransfers = transfers[:len(transfers)-1]
		found = true

		groupID := uuid.New()
		transfer.GroupID = uuid.NullUUID{UUID: groupID, Valid: true}
		openTransfer.GroupID = uuid.NullUUID{UUID: groupID, Valid: true}

		_, err = ynab.queries.CreateTransaction(context, transfer.CreateTransactionParams)
		if err != nil {
			return fmt.Errorf("save transaction %v: %w", transfer.CreateTransactionParams, err)
		}
		_, err = ynab.queries.CreateTransaction(context, openTransfer.CreateTransactionParams)
		if err != nil {
			return fmt.Errorf("save transaction %v: %w", openTransfer.CreateTransactionParams, err)
		}
		break
	}

	if !found {
		*openTransfers = append(*openTransfers, transfer)
	}
	return nil
}

func trimLastChar(s string) string {
	r, size := utf8.DecodeLastRuneInString(s)
	if r == utf8.RuneError && (size == 0 || size == 1) {
		size = 0
	}
	return s[:len(s)-size]
}

func GetAmount(inflow string, outflow string) (Numeric, error) {
	// Remove trailing currency
	inflow = strings.Replace(trimLastChar(inflow), ",", ".", 1)
	outflow = strings.Replace(trimLastChar(outflow), ",", ".", 1)

	num := Numeric{}
	err := num.Set(inflow)
	if err != nil {
		return num, fmt.Errorf("parse inflow %s: %w", inflow, err)
	}

	// if inflow is zero, use outflow
	if num.Int.Int64() != 0 {
		return num, nil
	}

	err = num.Set("-" + outflow)
	if err != nil {
		return num, fmt.Errorf("parse outflow %s: %w", inflow, err)
	}
	return num, nil
}

func (ynab *YNABImport) GetAccount(context context.Context, name string) (*Account, error) {
	for _, acc := range ynab.accounts {
		if acc.Name == name {
			return &acc, nil
		}
	}

	account, err := ynab.queries.CreateAccount(context, CreateAccountParams{Name: name, BudgetID: ynab.budgetID})
	if err != nil {
		return nil, err
	}

	ynab.accounts = append(ynab.accounts, account)
	return &account, nil
}

func (ynab *YNABImport) GetPayee(context context.Context, name string) (uuid.NullUUID, error) {
	if name == "" {
		return uuid.NullUUID{}, nil
	}

	for _, pay := range ynab.payees {
		if pay.Name == name {
			return uuid.NullUUID{UUID: pay.ID, Valid: true}, nil
		}
	}

	payee, err := ynab.queries.CreatePayee(context, CreatePayeeParams{Name: name, BudgetID: ynab.budgetID})
	if err != nil {
		return uuid.NullUUID{}, err
	}

	ynab.payees = append(ynab.payees, payee)
	return uuid.NullUUID{UUID: payee.ID, Valid: true}, nil
}

func (ynab *YNABImport) GetCategory(context context.Context, group string, name string) (uuid.NullUUID, error) { //nolint
	if group == "" || name == "" {
		return uuid.NullUUID{}, nil
	}

	for _, category := range ynab.categories {
		if category.Name == name && category.Group == group {
			return uuid.NullUUID{UUID: category.ID, Valid: true}, nil
		}
	}

	var categoryGroup CategoryGroup
	for _, existingGroup := range ynab.categoryGroups {
		if existingGroup.Name == group {
			categoryGroup = existingGroup
		}
	}

	if categoryGroup.Name == "" {
		newGroup := CreateCategoryGroupParams{Name: group, BudgetID: ynab.budgetID}
		var err error
		categoryGroup, err = ynab.queries.CreateCategoryGroup(context, newGroup)
		if err != nil {
			return uuid.NullUUID{}, err
		}
		ynab.categoryGroups = append(ynab.categoryGroups, categoryGroup)
	}

	newCategory := CreateCategoryParams{Name: name, CategoryGroupID: categoryGroup.ID}
	category, err := ynab.queries.CreateCategory(context, newCategory)
	if err != nil {
		return uuid.NullUUID{}, err
	}

	getCategory := GetCategoriesRow{
		ID:              category.ID,
		CategoryGroupID: category.CategoryGroupID,
		Name:            category.Name,
		Group:           categoryGroup.Name,
	}
	ynab.categories = append(ynab.categories, getCategory)
	return uuid.NullUUID{UUID: category.ID, Valid: true}, nil
}