-- Robust Leave Balance Trigger to handle ALL status transitions and edits
-- This handles Reset to Pending, Date Changes, and Role Changes

CREATE OR REPLACE FUNCTION public.update_leave_balance_on_status_change()
RETURNS TRIGGER AS $$
DECLARE
    v_is_deductible_old BOOLEAN;
    v_is_deductible_new BOOLEAN;
    v_year_old INTEGER;
    v_year_new INTEGER;
BEGIN
    -- 1. REVERT OLD STATE
    IF OLD IS NOT NULL THEN
        SELECT is_deductible INTO v_is_deductible_old FROM public.leave_types WHERE id = OLD.leave_type_id;
        
        IF v_is_deductible_old THEN
            v_year_old := EXTRACT(YEAR FROM OLD.start_date);
            
            -- Revert based on OLD status
            IF OLD.status = 'approved' THEN
                UPDATE public.leave_balances 
                SET used_days = used_days - OLD.total_days
                WHERE employee_id = OLD.employee_id AND leave_type_id = OLD.leave_type_id AND year = v_year_old;
            ELSIF OLD.status = 'pending' THEN
                UPDATE public.leave_balances 
                SET pending_days = pending_days - OLD.total_days
                WHERE employee_id = OLD.employee_id AND leave_type_id = OLD.leave_type_id AND year = v_year_old;
            END IF;
        END IF;
    END IF;

    -- 2. APPLY NEW STATE
    IF (TG_OP = 'INSERT' OR TG_OP = 'UPDATE') THEN
        IF NEW.status != 'cancelled' AND NEW.status != 'rejected' THEN
            SELECT is_deductible INTO v_is_deductible_new FROM public.leave_types WHERE id = NEW.leave_type_id;
            
            IF v_is_deductible_new THEN
                v_year_new := EXTRACT(YEAR FROM NEW.start_date);
                
                -- Ensure balance record exists
                INSERT INTO public.leave_balances (employee_id, leave_type_id, year)
                VALUES (NEW.employee_id, NEW.leave_type_id, v_year_new)
                ON CONFLICT (employee_id, leave_type_id, year) DO NOTHING;
                
                -- Apply based on NEW status
                IF NEW.status = 'approved' THEN
                    UPDATE public.leave_balances 
                    SET used_days = used_days + NEW.total_days
                    WHERE employee_id = NEW.employee_id AND leave_type_id = NEW.leave_type_id AND year = v_year_new;
                ELSIF NEW.status = 'pending' THEN
                    UPDATE public.leave_balances 
                    SET pending_days = pending_days + NEW.total_days
                    WHERE employee_id = NEW.employee_id AND leave_type_id = NEW.leave_type_id AND year = v_year_new;
                END IF;
            END IF;
        END IF;
    END IF;
    
    IF (TG_OP = 'DELETE') THEN
        RETURN OLD;
    END IF;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

DROP TRIGGER IF EXISTS tr_update_leave_balance ON public.leave_requests;
CREATE TRIGGER tr_update_leave_balance
AFTER INSERT OR UPDATE OR DELETE ON public.leave_requests
FOR EACH ROW EXECUTE PROCEDURE public.update_leave_balance_on_status_change();