Can you help me getting my buttons to work in Pygame?

Question

So I'm working on a snake game and I ran into some problems with the button class. I would like to implement two buttons at the end of the game. One in order to restart the game, the other one to close the game. I tested the buttons and they should be working. I just don't know where to draw them, because when I draw them at the end, I cannot press the button, because the game freezes due to the game.sleep() command. The game closes on itself, that's why I added a delay.

Pasting the whole code would be too much. I edited the code, so only the button code is shown below with the gameloop. I drew the buttons al the way at the end.

clicked = False

font = pygame.font.SysFont('Constantia', 30)



class button():

    # colours for button and text
    button_col = (255, 0, 0)
    hover_col = (75, 225, 255)
    click_col = (50, 150, 255)
    text_col = BLACK
    width = 180
    height = 70

    def __init__(self, x, y, text):
        self.x = x
        self.y = y
        self.text = text

    def draw_button(self):

        global clicked
        action = False

        # get mouse position
        pos = pygame.mouse.get_pos()

        # create pygame Rect object for the button
        button_rect = Rect(self.x, self.y, self.width, self.height)

        # check mouseover and clicked conditions
        if button_rect.collidepoint(pos):
            if pygame.mouse.get_pressed()[0] == 1:
                clicked = True
                pygame.draw.rect(SCREEN, self.click_col, button_rect)
            elif pygame.mouse.get_pressed()[0] == 0 and clicked == True:
                clicked = False
                action = True
            else:
                pygame.draw.rect(SCREEN, self.hover_col, button_rect)
        else:
            pygame.draw.rect(SCREEN, self.button_col, button_rect)

        # add shading to button
        pygame.draw.line(SCREEN, WHITE, (self.x, self.y),
                         (self.x + self.width, self.y), 2)
        pygame.draw.line(SCREEN, WHITE, (self.x, self.y),
                         (self.x, self.y + self.height), 2)
        pygame.draw.line(SCREEN, BLACK, (self.x, self.y + self.height),
                         (self.x + self.width, self.y + self.height), 2)
        pygame.draw.line(SCREEN, BLACK, (self.x + self.width, self.y),
                         (self.x + self.width, self.y + self.height), 2)

        # add text to button
        text_img = font.render(self.text, True, self.text_col)
        text_len = text_img.get_width()
        SCREEN.blit(text_img, (self.x + int(self.width / 2) -
                    int(text_len / 2), self.y + 25))
        return action


again = button(75, 200, 'Play Again?')
quit = button(325, 200, 'Quit?')


def main():

    RUN = True

    SNAKE_POS_X = BLOCKSIZE
    SNAKE_POS_Y = BLOCKSIZE
    SNAKE_POS_X_CHANGE = 0
    SNAKE_POS_Y_CHANGE = 0
    LENGTH_OF_SNAKE = 1

    global FOOD_POS_X, FOOD_POS_Y
    FOOD_POS_X = round(random.randrange(
        0, WIDTH - BLOCKSIZE) / BLOCKSIZE) * BLOCKSIZE
    FOOD_POS_Y = round(random.randrange(
        0, HEIGHT - BLOCKSIZE) / BLOCKSIZE) * BLOCKSIZE

    while RUN:

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                RUN = False

            # snake_movement
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP:
                    SNAKE_POS_X_CHANGE = 0
                    SNAKE_POS_Y_CHANGE = -BLOCKSIZE
                elif event.key == pygame.K_DOWN:
                    SNAKE_POS_X_CHANGE = 0
                    SNAKE_POS_Y_CHANGE = BLOCKSIZE
                elif event.key == pygame.K_RIGHT:
                    SNAKE_POS_X_CHANGE = BLOCKSIZE
                    SNAKE_POS_Y_CHANGE = 0
                elif event.key == pygame.K_LEFT:
                    SNAKE_POS_X_CHANGE = -BLOCKSIZE
                    SNAKE_POS_Y_CHANGE = 0

        
        if SNAKE_POS_X >= WIDTH or SNAKE_POS_X < 0 or SNAKE_POS_Y >= HEIGHT or SNAKE_POS_Y < 0:
            RUN = False

        SNAKE_POS_X += SNAKE_POS_X_CHANGE
        SNAKE_POS_Y += SNAKE_POS_Y_CHANGE

        SCREEN.fill(BISQUE2)
        checkerboard()
        food()
        SNAKE_HEAD = []
        SNAKE_HEAD.append(SNAKE_POS_X)
        SNAKE_HEAD.append(SNAKE_POS_Y)
        SNAKE_LIST.append(SNAKE_HEAD)
        if len(SNAKE_LIST) > LENGTH_OF_SNAKE:
            del SNAKE_LIST[0]

        for x in SNAKE_LIST[:-1]:
            if x == SNAKE_HEAD:
                RUN = False

        snake(BLOCKSIZE, SNAKE_LIST)
        score(LENGTH_OF_SNAKE - 1)
        # draw_grid()
        CLOCK.tick(FPS)

        pygame.display.update()

        if SNAKE_POS_X == FOOD_POS_X and SNAKE_POS_Y == FOOD_POS_Y:
            FOOD_POS_X = round(random.randrange(
                0, WIDTH - BLOCKSIZE) / BLOCKSIZE) * BLOCKSIZE
            FOOD_POS_Y = round(random.randrange(
                0, HEIGHT - BLOCKSIZE) / BLOCKSIZE) * BLOCKSIZE
            LENGTH_OF_SNAKE += 1
            CRUNCH.play()

    game_over_message("Game Over!", BLACK)
    GAME_OVER_SOUND.play()
    # pygame.display.update()

    if again.draw_button():
        main()

    if quit.draw_button():
        pygame.quit()

    pygame.display.update()

    time.sleep(2)
    pygame.quit()


main()

You can find my code on Github: https://github.com/Lin3x10/Snake-MA/blob/master/snake.py