Personal-Finance-Management/app/routes/budgets.py

from flask import Blueprint, render_template, redirect, url_for, flash, request, jsonify
from flask_login import login_required
from flask_wtf import FlaskForm
from wtforms import SelectField, DecimalField, BooleanField, SubmitField, HiddenField
from wtforms.validators import DataRequired, NumberRange
from app.extensions import db
from app.models.budget import Budget
from app.models.category import Category
from app.services.budget_service import get_budget_summary, get_total_budget, get_total_spent
from datetime import date

budgets_bp = Blueprint('budgets', __name__, url_prefix='/budgets')


def _current_month():
    return date.today().strftime('%Y-%m')


def _expense_category_choices():
    cats = Category.query.filter(
        Category.category_type.in_(['expense', 'both']),
        Category.is_active == True,
        Category.parent_id == None,
    ).order_by(Category.name).all()
    return [(str(c.id), c.name) for c in cats]


class BudgetForm(FlaskForm):
    category_id = SelectField('Category', validators=[DataRequired()])
    month = HiddenField()
    limit_amount = DecimalField('Monthly Limit', validators=[DataRequired(), NumberRange(min=0.01)], places=2)
    rollover_enabled = BooleanField('Roll over unused amount to next month')
    submit = SubmitField('Save Budget')


@budgets_bp.route('/')
@login_required
def index():
    month = request.args.get('month', _current_month())
    summary = get_budget_summary(month)
    total_budget = get_total_budget(month)
    total_spent = get_total_spent(month)
    overall_pct = min(round((total_spent / total_budget * 100), 1), 100) if total_budget > 0 else 0

    # Month navigation
    year, mo = map(int, month.split('-'))
    if mo == 1:
        prev_month = f'{year-1}-12'
    else:
        prev_month = f'{year}-{mo-1:02d}'
    if mo == 12:
        next_month = f'{year+1}-01'
    else:
        next_month = f'{year}-{mo+1:02d}'

    return render_template('budgets/index.html',
                           month=month,
                           month_label=date(year, mo, 1).strftime('%B %Y'),
                           prev_month=prev_month,
                           next_month=next_month,
                           summary=summary,
                           total_budget=total_budget,
                           total_spent=total_spent,
                           overall_pct=overall_pct)


@budgets_bp.route('/new', methods=['GET', 'POST'])
@login_required
def new():
    month = request.args.get('month', _current_month())
    form = BudgetForm()
    form.category_id.choices = _expense_category_choices()
    form.month.data = month

    if form.validate_on_submit():
        month_val = form.month.data or month
        existing = Budget.query.filter_by(
            category_id=int(form.category_id.data),
            month=month_val
        ).first()
        if existing:
            flash('Budget for this category already exists this month. Edit it instead.', 'warning')
            return redirect(url_for('budgets.index', month=month_val))

        budget = Budget(
            category_id=int(form.category_id.data),
            month=month_val,
            limit_amount=form.limit_amount.data,
            rollover_enabled=form.rollover_enabled.data,
            rollover_amount=0,
        )
        db.session.add(budget)
        db.session.commit()
        flash('Budget created.', 'success')
        return redirect(url_for('budgets.index', month=month_val))

    return render_template('budgets/form.html', form=form, month=month,
                           title='New Budget', month_label=date(*map(int, month.split('-')), 1).strftime('%B %Y'))


@budgets_bp.route('/<int:id>/edit', methods=['GET', 'POST'])
@login_required
def edit(id):
    budget = db.get_or_404(Budget, id)
    form = BudgetForm(obj=budget)
    form.category_id.choices = _expense_category_choices()

    if request.method == 'GET':
        form.category_id.data = str(budget.category_id)
        form.month.data = budget.month

    if form.validate_on_submit():
        budget.limit_amount = form.limit_amount.data
        budget.rollover_enabled = form.rollover_enabled.data
        db.session.commit()
        flash('Budget updated.', 'success')
        return redirect(url_for('budgets.index', month=budget.month))

    year, mo = map(int, budget.month.split('-'))
    return render_template('budgets/form.html', form=form, month=budget.month,
                           title='Edit Budget', budget=budget,
                           month_label=date(year, mo, 1).strftime('%B %Y'))


@budgets_bp.route('/<int:id>/delete', methods=['POST'])
@login_required
def delete(id):
    budget = db.get_or_404(Budget, id)
    month = budget.month
    db.session.delete(budget)
    db.session.commit()
    flash('Budget removed.', 'info')
    return redirect(url_for('budgets.index', month=month))


@budgets_bp.route('/copy', methods=['POST'])
@login_required
def copy_month():
    """Copy all budgets from one month to another."""
    from_month = request.form.get('from_month')
    to_month = request.form.get('to_month')
    if not from_month or not to_month:
        flash('Invalid months.', 'danger')
        return redirect(url_for('budgets.index'))

    from_budgets = Budget.query.filter_by(month=from_month).all()
    copied = 0
    for fb in from_budgets:
        existing = Budget.query.filter_by(
            month=to_month, category_id=fb.category_id
        ).first()
        if not existing:
            db.session.add(Budget(
                category_id=fb.category_id,
                month=to_month,
                limit_amount=fb.limit_amount,
                rollover_enabled=fb.rollover_enabled,
                rollover_amount=0,
            ))
            copied += 1
    db.session.commit()
    flash(f'Copied {copied} budget(s) to {to_month}.', 'success')
    return redirect(url_for('budgets.index', month=to_month))