Novault-backend/internal/repository/budget_repository.go

package repository

import (
	"errors"
	"fmt"
	"time"

	"accounting-app/internal/models"

	"gorm.io/gorm"
)

// Common repository errors
var (
	ErrBudgetNotFound = errors.New("budget not found")
	ErrBudgetInUse    = errors.New("budget is in use and cannot be deleted")
)

// BudgetRepository handles database operations for budgets
type BudgetRepository struct {
	db *gorm.DB
}

// NewBudgetRepository creates a new BudgetRepository instance
func NewBudgetRepository(db *gorm.DB) *BudgetRepository {
	return &BudgetRepository{db: db}
}

// Create creates a new budget in the database
func (r *BudgetRepository) Create(budget *models.Budget) error {
	if err := r.db.Create(budget).Error; err != nil {
		return fmt.Errorf("failed to create budget: %w", err)
	}
	return nil
}

// GetByID retrieves a budget by its ID
func (r *BudgetRepository) GetByID(userID uint, id uint) (*models.Budget, error) {
	var budget models.Budget
	if err := r.db.Preload("Category").Preload("Account").Where("user_id = ?", userID).First(&budget, id).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return nil, ErrBudgetNotFound
		}
		return nil, fmt.Errorf("failed to get budget: %w", err)
	}
	return &budget, nil
}

// GetAll retrieves all budgets for a user
func (r *BudgetRepository) GetAll(userID uint) ([]models.Budget, error) {
	var budgets []models.Budget
	if err := r.db.Preload("Category").Preload("Account").Where("user_id = ?", userID).Order("created_at DESC").Find(&budgets).Error; err != nil {
		return nil, fmt.Errorf("failed to get budgets: %w", err)
	}
	return budgets, nil
}

// Update updates an existing budget in the database
func (r *BudgetRepository) Update(budget *models.Budget) error {
	// First check if the budget exists
	var existing models.Budget
	if err := r.db.Where("user_id = ?", budget.UserID).First(&existing, budget.ID).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return ErrBudgetNotFound
		}
		return fmt.Errorf("failed to check budget existence: %w", err)
	}

	// Update the budget
	if err := r.db.Save(budget).Error; err != nil {
		return fmt.Errorf("failed to update budget: %w", err)
	}
	return nil
}

// Delete deletes a budget by its ID
func (r *BudgetRepository) Delete(userID uint, id uint) error {
	// First check if the budget exists
	var budget models.Budget
	if err := r.db.Where("user_id = ?", userID).First(&budget, id).Error; err != nil {
		if errors.Is(err, gorm.ErrRecordNotFound) {
			return ErrBudgetNotFound
		}
		return fmt.Errorf("failed to check budget existence: %w", err)
	}

	// Delete the budget (soft delete due to gorm.DeletedAt field)
	if err := r.db.Delete(&budget).Error; err != nil {
		return fmt.Errorf("failed to delete budget: %w", err)
	}
	return nil
}

// GetByCategoryID retrieves all budgets for a specific category and user
func (r *BudgetRepository) GetByCategoryID(userID, categoryID uint) ([]models.Budget, error) {
	var budgets []models.Budget
	if err := r.db.Preload("Category").Preload("Account").Where("user_id = ? AND category_id = ?", userID, categoryID).Order("created_at DESC").Find(&budgets).Error; err != nil {
		return nil, fmt.Errorf("failed to get budgets by category: %w", err)
	}
	return budgets, nil
}

// GetByAccountID retrieves all budgets for a specific account and user
func (r *BudgetRepository) GetByAccountID(userID, accountID uint) ([]models.Budget, error) {
	var budgets []models.Budget
	if err := r.db.Preload("Category").Preload("Account").Where("user_id = ? AND account_id = ?", userID, accountID).Order("created_at DESC").Find(&budgets).Error; err != nil {
		return nil, fmt.Errorf("failed to get budgets by account: %w", err)
	}
	return budgets, nil
}

// GetByPeriodType retrieves all budgets of a specific period type for a user
func (r *BudgetRepository) GetByPeriodType(userID uint, periodType models.PeriodType) ([]models.Budget, error) {
	var budgets []models.Budget
	if err := r.db.Preload("Category").Preload("Account").Where("user_id = ? AND period_type = ?", userID, periodType).Order("created_at DESC").Find(&budgets).Error; err != nil {
		return nil, fmt.Errorf("failed to get budgets by period type: %w", err)
	}
	return budgets, nil
}

// GetActiveBudgets retrieves all budgets that are currently active for a user
func (r *BudgetRepository) GetActiveBudgets(userID uint, currentDate time.Time) ([]models.Budget, error) {
	var budgets []models.Budget
	query := r.db.Preload("Category").Preload("Account").
		Where("user_id = ?", userID).
		Where("start_date <= ?", currentDate).
		Where("end_date IS NULL OR end_date >= ?", currentDate).
		Order("created_at DESC")

	if err := query.Find(&budgets).Error; err != nil {
		return nil, fmt.Errorf("failed to get active budgets: %w", err)
	}
	return budgets, nil
}

// GetSpentAmount calculates the total spent amount for a budget within a specific date range
func (r *BudgetRepository) GetSpentAmount(budget *models.Budget, startDate, endDate time.Time) (float64, error) {
	var totalSpent float64

	query := r.db.Model(&models.Transaction{}).
		Where("user_id = ?", budget.UserID).
		Where("type = ?", models.TransactionTypeExpense).
		Where("transaction_date >= ? AND transaction_date <= ?", startDate, endDate)

	// Filter by category if specified
	if budget.CategoryID != nil {
		query = query.Where("category_id = ?", *budget.CategoryID)
	}

	// Filter by account if specified
	if budget.AccountID != nil {
		query = query.Where("account_id = ?", *budget.AccountID)
	}

	if err := query.Select("COALESCE(SUM(amount), 0)").Scan(&totalSpent).Error; err != nil {
		return 0, fmt.Errorf("failed to calculate spent amount: %w", err)
	}

	return totalSpent, nil
}

// ExistsByID checks if a budget with the given ID exists for a user
func (r *BudgetRepository) ExistsByID(userID, id uint) (bool, error) {
	var count int64
	if err := r.db.Model(&models.Budget{}).Where("user_id = ? AND id = ?", userID, id).Count(&count).Error; err != nil {
		return false, fmt.Errorf("failed to check budget existence: %w", err)
	}
	return count > 0, nil
}