import { StateCreator } from 'zustand'
import { MonthlyBudget } from '@/lib/types'

const now = new Date()

export interface BudgetSlice {
    monthlyBudgets: MonthlyBudget[]
    currentMonth: number
    currentYear: number

    setMonthlyBudget: (budget: MonthlyBudget) => void
    updateMonthlyBudget: (month: number, year: number, updates: Partial<MonthlyBudget>) => void
}

export const createBudgetSlice: StateCreator<BudgetSlice> = (set) => ({
    monthlyBudgets: [],
    currentMonth: now.getMonth() + 1,
    currentYear: now.getFullYear(),

    setMonthlyBudget: (budget) =>
        set((state) => {
            const existingIndex = state.monthlyBudgets.findIndex(
                (b) => b.month === budget.month && b.year === budget.year
            )
            if (existingIndex >= 0) {
                const newBudgets = [...state.monthlyBudgets]
                newBudgets[existingIndex] = budget
                return { monthlyBudgets: newBudgets }
            }
            return { monthlyBudgets: [...state.monthlyBudgets, budget] }
        }),

    updateMonthlyBudget: (month, year, updates) =>
        set((state) => ({
            monthlyBudgets: state.monthlyBudgets.map((b) =>
                b.month === month && b.year === year ? { ...b, ...updates } : b
            ),
        })),
})