firefly_budget_overview.py

"""
Firefly III Budget Overview Generator

This script generates a budget overview using data from a Firefly III instance.
It fetches salary information, budget details, and transactions, then calculates
and displays budget overviews for each pay cycle.

The script uses environment variables for API configuration and budget names.
It supports various output formats and configurable logging levels.

Usage:
    python firefly_budget_overview.py [--no-lines] [--mobile] [--log-level {debug,info,warning,error,critical}]

Options:
    --no-lines  Remove column and row lines from the output
    --mobile    Output in mobile-friendly format
    --log-level Set the logging level (debug, info, warning, error, critical)

Environment Variables:
    FIREFLY_API_URL: URL of the Firefly III API
    FIREFLY_API_TOKEN: API token for authentication
    FIREFLY_BUDGETS: Comma-separated list of budget names to include in the overview

# Note: The above need to be specified in a .env file located in the same directory as this script.
# Environment variables example:
FIREFLY_API_URL=http://192.168.1.10:8080/api/v1
FIREFLY_API_TOKEN=<Personal_Access_Token>
FIREFLY_BUDGETS=Car Fuel,Eating Out,Extra Expenses,Groceries

Firefly requirements:
    - The budgets setup in Firefly III.
    - The relevant transactions to be counted are associated with the budgets.
    - The salary transactions are categorized as "Salary".
    - Salary transactions are expected to be monthly.
"""

import os
import requests
import pandas as pd
from datetime import datetime, timedelta
import argparse
import logging
from tabulate import tabulate
from dotenv import load_dotenv
from typing import Any, TypedDict
from dateutil.relativedelta import relativedelta

# Set up logging
logger = logging.getLogger(__name__)

# Load environment variables from .env file
_ = load_dotenv()

# Firefly III API configuration
API_URL = os.getenv('FIREFLY_API_URL')
API_TOKEN = os.getenv('FIREFLY_API_TOKEN')

HEADERS = {
    "Authorization": f"Bearer {API_TOKEN}",
    "Accept": "application/json",
    "Content-Type": "application/json"
}

# Get budgets from environment variable
BUDGETS = os.getenv('FIREFLY_BUDGETS', '').split(',')

class FireflyTransaction(TypedDict):
    date: str
    amount: str
    budget_name: str

def get_firefly_data(endpoint: str, params: dict[str, Any] | None = None) -> dict[str, list[dict[str, Any]] | dict[str, Any]]:
    """
    Fetch data from the Firefly III API.

    Args:
        endpoint (str): The API endpoint to query.
        params (dict[str, Any] | None): Optional query parameters.

    Returns:
        dict[str, list[dict[str, Any]] | dict[str, Any]]: The API response data.
    """
    all_data: list[dict[str, Any]] = []
    page = 1
    while True:
        if params is None:
            params = {}
        params['page'] = page
        response = requests.get(f"{API_URL}/{endpoint}", headers=HEADERS, params=params)
        response.raise_for_status()
        data: dict[str, Any] | list[dict[str, Any]] = response.json()
        if isinstance(data, list):
            return {"data": data}
        if isinstance(data, dict) and 'data' in data:
            if isinstance(data['data'], list):
                all_data.extend(data['data'])
            else:
                return data  # Return the entire response if 'data' is not a list
        if isinstance(data, dict) and 'links' in data and 'next' in data['links'] and data['links']['next']:
            page += 1
        else:
            break
    return {'data': all_data}

def get_salaries(start_date: datetime, end_date: datetime) -> pd.DataFrame:
    """
    Fetch and process salary transactions within the given date range.

    Args:
        start_date (datetime): The start date for fetching transactions.
        end_date (datetime): The end date for fetching transactions.

    Returns:
        pd.DataFrame: A DataFrame containing processed salary information.
    """
    logger.info(f"Getting salaries with {start_date} start date and {end_date} end date...")
    transactions = get_firefly_data("transactions", {
        "start": start_date.strftime("%Y-%m-%d"),
        "end": end_date.strftime("%Y-%m-%d"),
        "type": "deposit",
        "category": "Salary"
    })
    
    if not transactions['data']:
        logger.warning("No salary transactions found.")
        return pd.DataFrame()

    salaries = []

    for t in transactions['data']:
        try:
            transaction = t['attributes']['transactions'][0]
            date = pd.to_datetime(transaction['date'])
            amount = float(transaction['amount'])
            category = transaction.get('category_name', '')
            
            if category == 'Salary':
                salaries.append({
                    'Date': date,
                    'Amount': amount
                })
        except (KeyError, IndexError) as e:
            logger.error(f"Unexpected transaction structure: {e}")
            logger.debug(f"Transaction data: {t}")
    
    if not salaries:
        logger.warning("No salary transactions found after processing.")
        return pd.DataFrame()
    
    df = pd.DataFrame(salaries).sort_values('Date')
    df = df.drop_duplicates(subset=['Date'], keep='first')
    
    logger.info(f"Found {len(df)} unique salary transactions:")
    logger.debug(f"Salary transactions:\n{df}")
    
    return df

def get_budget_info() -> pd.DataFrame:
    """
    Fetch and process budget information for the specified budgets.

    Returns:
        pd.DataFrame: A DataFrame containing budget information.
    """
    logger.info("Getting budgets info...")
    budgets = get_firefly_data("budgets")
    
    budget_info = []
    for budget in budgets['data']:
        if budget['attributes']['name'] in BUDGETS:
            budget_id = budget['id']
            budget_details = get_firefly_data(f"budgets/{budget_id}")
            
            # Check if 'data' is a list or a dictionary
            if isinstance(budget_details['data'], list):
                attributes = budget_details['data'][0]['attributes']
            else:
                attributes = budget_details['data']['attributes']
            
            budget_info.append({
                'Category': attributes['name'],
                'Budget Amount': float(attributes['auto_budget_amount'] or 0),
                'Auto Budget Type': attributes['auto_budget_type'],
                'Auto Budget Period': attributes['auto_budget_period']
            })
    
    df = pd.DataFrame(budget_info)
    logger.info("Final budget info:")
    logger.debug(f"Budget info:\n{df}")
    return df

def get_transactions(start_date: datetime, end_date: datetime) -> pd.DataFrame:
    """
    Fetch and process transactions within the given date range.

    Args:
        start_date (datetime): The start date for fetching transactions.
        end_date (datetime): The end date for fetching transactions.

    Returns:
        pd.DataFrame: A DataFrame containing processed transaction information.
    """
    logger.info("Getting transactions...")
    transactions = get_firefly_data("transactions", {
        "start": start_date.strftime("%Y-%m-%d"),
        "end": end_date.strftime("%Y-%m-%d"),
        "type": "withdrawal"
    })
    
    if not transactions['data']:
        logger.warning("No transactions found.")
        return pd.DataFrame()

    df = pd.DataFrame([
        {
            'Date': pd.to_datetime(t['attributes']['transactions'][0]['date']),
            'Amount': float(t['attributes']['transactions'][0]['amount']),
            'Category': t['attributes']['transactions'][0].get('budget_name', 'No Budget')
        }
        for t in transactions['data']
    ])
    df = df[df['Category'].isin(BUDGETS)]  # Filter to include only specified budgets
    logger.info("Transactions DataFrame:")
    logger.debug(f"Transactions:\n{df}")
    return df

def calculate_spent_amount(transactions: pd.DataFrame, category: str, start_date: pd.Timestamp, end_date: pd.Timestamp) -> float:
    """
    Calculate the total amount spent for a specific category within a date range.

    Args:
        transactions (pd.DataFrame): The DataFrame containing all transactions.
        category (str): The budget category to calculate for.
        start_date (pd.Timestamp): The start date of the calculation period.
        end_date (pd.Timestamp): The end date of the calculation period.

    Returns:
        float: The total amount spent, rounded to 3 decimal places.
    """

    logger.debug(f"Calculating spent amount for {category} category with {start_date} start date and {end_date} end date...")
    filtered_transactions = transactions[
        (transactions['Category'] == category) &
        (transactions['Date'] >= start_date) &
        (transactions['Date'] < end_date)
    ]
    spent = round(filtered_transactions['Amount'].sum(), 3)  # Round to 3 decimal places
    logger.debug(f"Category {category}, Start: {start_date.date()}, End: {end_date.date()}")
    logger.debug(f"Filtered transactions:\n{filtered_transactions}")
    logger.debug(f"Spent amount: {spent}")
    return spent

def get_budget_period(pay_date: pd.Timestamp, next_pay_date: pd.Timestamp) -> tuple[pd.Timestamp, pd.Timestamp]:
    """
    Generate a comprehensive budget overview.

    This function orchestrates the entire process of fetching data,
    calculating budgets, and preparing the overview for display.

    Args:
        show_lines (bool): Whether to show lines in the table output.
        mobile (bool): Whether to format the output for mobile display.

    Returns:
        list[dict[str, Any]] | None: A list of budget overview data, or None if data couldn't be generated.
    """
    logger.debug(f"Getting budget period for {pay_date} pay date with {next_pay_date} next pay date...")
    period_end = next_pay_date - pd.Timedelta(days=1)
    if period_end < pay_date:
        logger.warning(f"Calculated end date {period_end} is before pay date {pay_date}. Using next_pay_date as end.")
        period_end = next_pay_date
    return pay_date, period_end

def get_pay_cycles(salaries: pd.DataFrame) -> list[dict[str, Any]]:
    """
    Generate pay cycles based on salary data.

    This function takes a DataFrame of salary information and creates a list of pay cycles.
    Each pay cycle includes the start and end dates, as well as formatted strings for display.

    Args:
        salaries (pd.DataFrame): A DataFrame containing salary information with 'Date' column.

    Returns:
        list[dict[str, Any]]: A list of dictionaries, each representing a pay cycle.
        Each dictionary contains 'Pay Cycle', 'Budget Period', 'Start', and 'End' keys.
    """

    logger.info("Getting pay cycles...")
    pay_cycles: list[dict[str, Any]] = []
    for i in range(len(salaries)):
        salary_row = salaries.iloc[i]
        pay_date = pd.Timestamp(salary_row['Date'])
        logger.debug(f"Pay date: {pay_date}")
        if i + 1 < len(salaries):
            next_salary_date = pd.Timestamp(salaries.iloc[i+1]['Date'])
            logger.debug(f"Next pay date with another salary: {next_salary_date}")
        else:
            # For the last salary, set the next pay date to one month after
            next_salary_date = pay_date + relativedelta(months=1)
            logger.debug(f"Last salary. Next pay date set to one month later: {next_salary_date}")
        
        period_start, period_end = get_budget_period(pay_date, next_salary_date)
        pay_cycles.append({
            'Pay Cycle': pay_date.date().isoformat(),
            'Budget Period': f'{period_start.date().isoformat()} to {period_end.date().isoformat()}',
            'Start': period_start,
            'End': period_end
        })
        logger.debug(f"Period start: {period_start}, Period end: {period_end}")
    return pay_cycles

def calculate_budget_overview(pay_cycle: dict[str, Any], budget_info: pd.DataFrame, transactions: pd.DataFrame) -> list[dict[str, Any]]:
    """
    Calculate the budget overview for a specific pay cycle.

    This function processes the budget information and transactions for a given pay cycle,
    calculating spent amounts, adjusting displays, and formatting differences.

    Args:
        pay_cycle (dict[str, Any]): A dictionary representing a pay cycle with 'Start' and 'End' dates.
        budget_info (pd.DataFrame): A DataFrame containing budget information.
        transactions (pd.DataFrame): A DataFrame containing transaction data.

    Returns:
        list[dict[str, Any]]: A list of dictionaries, each representing a budget category overview
        for the given pay cycle.
    """

    logger.info("Calculating budget overview...")
    cycle_overview = []
    for _, budget_row in budget_info.iterrows():
        category = str(budget_row['Category'])
        budget_amount = float(budget_row['Budget Amount'])
        auto_budget_period = str(budget_row['Auto Budget Period'])
        
        spent_amount = calculate_spent_amount(transactions, category, pay_cycle['Start'], pay_cycle['End'])
        
        display_budget, display_spent = adjust_budget_display(auto_budget_period, budget_amount, spent_amount, pay_cycle['Start'], pay_cycle['End'])
        
        difference = round(display_budget - display_spent, 3)  # Round difference to 3 decimal places
        difference_text = format_difference(difference)
        
        cycle_overview.append({
            'Category': category,
            'Budget Type': auto_budget_period.capitalize(),
            'Budget Amount': f'{display_budget:.3f} AUD',  # Display with 3 decimal places
            'Spent Amount': f'{display_spent:.3f} AUD',  # Display with 3 decimal places
            'Difference': difference_text
        })
    return cycle_overview

def adjust_budget_display(auto_budget_period: str, budget_amount: float, spent_amount: float, start_date: pd.Timestamp, end_date: pd.Timestamp) -> tuple[float, float]:
    """
    Adjust the budget and spent amounts for display.

    This function adjusts the budget and spent amounts based on the budget period type.
    For weekly budgets, it calculates a monthly equivalent.

    Args:
        auto_budget_period (str): The type of budget period ('weekly' or other).
        budget_amount (float): The original budget amount.
        spent_amount (float): The amount spent.
        start_date (pd.Timestamp): The start date of the budget period.
        end_date (pd.Timestamp): The end date of the budget period.

    Returns:
        tuple[float, float]: A tuple containing the adjusted display budget and display spent amounts.
    """

    logger.debug("Adjusting budget display...")
    if auto_budget_period == 'weekly':
        monthly_equivalent = budget_amount * 52 / 12
        display_budget = round(monthly_equivalent, 3)  # Round to 3 decimal places
        display_spent = round(spent_amount * 30.44 / (end_date - start_date).days, 3)  # Round to 3 decimal places
    else:
        display_budget = round(budget_amount, 3)  # Round to 3 decimal places
        display_spent = round(spent_amount, 3)  # Round to 3 decimal places
    return display_budget, display_spent

def format_difference(difference: float) -> str:
    """
    Format the difference between budget and spent amounts.

    This function formats the difference as a string with color coding.
    Positive differences (under budget) are green, negative (over budget) are red.

    Args:
        difference (float): The difference between budget and spent amounts.

    Returns:
        str: A formatted string representing the difference, with ANSI color codes.
    """

    logger.debug("Formatting difference...")
    difference_text = f'{difference:.3f} AUD'  # Display with 3 decimal places
    if difference >= 0:
        return f'\033[92m{difference_text}\033[0m'  # Green text
    else:
        return f'\033[91m{difference_text}\033[0m'  # Red text

def generate_budget_overview(show_lines: bool, mobile: bool) -> list[dict[str, Any]] | None:
    """
    Generate a comprehensive budget overview for the last 90 days.

    This function orchestrates the entire process of generating a budget overview:
    1. Fetches salary data for the last 90 days
    2. Retrieves budget information
    3. Fetches transactions for the period
    4. Calculates pay cycles
    5. Generates a detailed budget overview for each pay cycle
    6. Displays the budget overview

    The function will return None if either salary data or budget information is not available.

    Args:
        show_lines (bool): If True, the displayed overview will include grid lines.
        mobile (bool): If True, the overview will be formatted for mobile display.

    Returns:
        list[dict[str, Any]] | None: A list of dictionaries, each representing a pay cycle with its budget overview.
        Each dictionary contains:
        - 'Pay Cycle': The start date of the pay cycle
        - 'Budget Period': The date range for the budget period
        - 'Start': The start date of the budget period
        - 'End': The end date of the budget period
        - 'Budgets': A list of budget overviews for each category in this pay cycle

        Returns None if salary data or budget information is not available.

    Logs:
        - Info level logs for major steps and found data counts
        - Warning level logs for missing data
        - Debug level logs for detailed transaction information and pay cycle processing

    Note:
        This function relies on several helper functions:
        - get_salaries()
        - get_budget_info()
        - get_transactions()
        - get_pay_cycles()
        - calculate_budget_overview()
        - display_budget_overview()
    """

    logger.info("Generating budget overview...")
    end_date = datetime.now()
    start_date = end_date - timedelta(days=90)  # Fetch last 90 days of data

    salaries = get_salaries(start_date, end_date)
    if salaries.empty:
        logger.warning("No salary data found. Cannot generate budget overview.")
        return None

    logger.info(f"Found {len(salaries)} salary transactions:")

    budget_info = get_budget_info()
    logger.info("\nBudget information:")
    logger.debug(f"Budget info:\n{budget_info}")

    if budget_info.empty:
        logger.warning("No budget information found. Cannot generate budget overview.")
        return None

    transactions = get_transactions(start_date, end_date)
    logger.info(f"\nFound {len(transactions)} budget-related transactions")
    logger.debug("Debug: All transactions:")
    logger.debug(f"{transactions}")

    pay_cycles = get_pay_cycles(salaries)

    budget_overview: list[dict[str, Any]] = []
    for pay_cycle in pay_cycles:
        logger.debug(f"\nDebug: Processing pay cycle '{pay_cycle['Pay Cycle']}' for '{pay_cycle['Budget Period']}' budget period...")
        cycle_overview = calculate_budget_overview(pay_cycle, budget_info, transactions)
        pay_cycle['Budgets'] = cycle_overview
        budget_overview.append(pay_cycle)

    display_budget_overview(budget_overview, show_lines, mobile)

    return budget_overview

def display_budget_overview(budget_overview: list[dict[str, Any]], show_lines: bool, mobile: bool) -> None:
    """
    Display the budget overview.

    This function chooses the appropriate display method based on the provided flags.

    Args:
        budget_overview (list[dict[str, Any]]): The complete budget overview data.
        show_lines (bool): If True, display with grid lines.
        mobile (bool): If True, use mobile-friendly format.

    Returns:
        None
    """

    logger.info("Displaying budget overview...")
    if mobile:
        display_mobile(budget_overview)
    elif show_lines:
        display_with_lines(budget_overview)
    else:
        display_without_lines(budget_overview)

def display_mobile(budget_overview: list[dict[str, Any]]) -> None:
    """
    Display the budget overview in a mobile-friendly format.

    This function prints the budget overview in a compact format suitable for mobile devices.

    Args:
        budget_overview (list[dict[str, Any]]): The complete budget overview data.

    Returns:
        None
    """

    logger.info("Displaying mobile budget overview...")
    for cycle in budget_overview:
        print(f"\nPay Cycle: {cycle['Pay Cycle']}")
        print(f"Budget Period: {cycle['Budget Period']}")
        for budget in cycle['Budgets']:
            print(f"  Category: {budget['Category']}")
            print(f"  Budget Type: {budget['Budget Type']}")
            print(f"  Budget Amount: {budget['Budget Amount']}")
            print(f"  Spent Amount: {budget['Spent Amount']}")
            print(f"  Difference: {budget['Difference']}")
            print('-' * 30)

def display_with_lines(budget_overview: list[dict[str, Any]]) -> None:
    """
    Display the budget overview with grid lines.

    This function prints the budget overview in a tabular format with grid lines.

    Args:
        budget_overview (list[dict[str, Any]]): The complete budget overview data.

    Returns:
        None
    """

    for cycle in budget_overview:
        print(f"\nPay Cycle: {cycle['Pay Cycle']}")
        print(f"Budget Period: {cycle['Budget Period']}")
        headers = ['Category', 'Budget Type', 'Budget Amount', 'Spent Amount', 'Difference']
        table = [[budget[header] for header in headers] for budget in cycle['Budgets']]
        print(tabulate(table, headers=headers, tablefmt="grid"))

def display_without_lines(budget_overview: list[dict[str, Any]]) -> None:
    """
    Display the budget overview without grid lines.

    This function prints the budget overview in a simple tabular format without grid lines.

    Args:
        budget_overview (list[dict[str, Any]]): The complete budget overview data.

    Returns:
        None
    """

    for cycle in budget_overview:
        print(f"\nPay Cycle: {cycle['Pay Cycle']}")
        print(f"Budget Period: {cycle['Budget Period']}")
        headers = ['Category', 'Budget Type', 'Budget Amount', 'Spent Amount', 'Difference']
        table = [[budget[header] for header in headers] for budget in cycle['Budgets']]
        print(tabulate(table, headers=headers, tablefmt="simple"))

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Generate a budget overview from Firefly III data.')
    parser.add_argument('--no-lines', action='store_true', help='Remove column and row lines from the output')
    parser.add_argument('--mobile', action='store_true', help='Output in mobile-friendly format')
    parser.add_argument('--log-level', choices=['debug', 'info', 'warning', 'error', 'critical'], 
                        help='Set the logging level')
    args = parser.parse_args()

    # Set up logging based on the log-level flag
    if args.log_level:
        numeric_level = getattr(logging, args.log_level.upper(), None)
        if not isinstance(numeric_level, int):
            raise ValueError(f'Invalid log level: {args.log_level}')
        logging.basicConfig(
            level=numeric_level,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
    else:
        # If no log level is specified, disable logging
        logging.disable(logging.CRITICAL)

    if args.log_level:
        logger.info(f"API URL: {API_URL}")
        logger.info(f"Budgets: {BUDGETS}")
        logger.info(f"Start date: {datetime.now() - timedelta(days=90)}")
        logger.info(f"End date: {datetime.now()}")

    _ = generate_budget_overview(not args.no_lines, args.mobile)