package middleware

import (
	"errors"

	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/fiber/v2/middleware/session"
)

const (
	flashSuccessKey = "flash_success"
	flashErrorKey   = "flash_error"
)

func SetFlashSuccess(c *fiber.Ctx, message string) error {
	return setFlash(c, flashSuccessKey, message)
}

func SetFlashError(c *fiber.Ctx, message string) error {
	return setFlash(c, flashErrorKey, message)
}

func ConsumeFlash() fiber.Handler {
	return func(c *fiber.Ctx) error {
		store, ok := c.Locals(contextStoreKey).(*session.Store)
		if !ok || store == nil {
			return errors.New("session store not available")
		}

		sess, err := store.Get(c)
		if err != nil {
			return err
		}

		success, _ := sess.Get(flashSuccessKey).(string)
		errMsg, _ := sess.Get(flashErrorKey).(string)

		sess.Delete(flashSuccessKey)
		sess.Delete(flashErrorKey)
		if err := sess.Save(); err != nil {
			return err
		}

		if success != "" {
			setTemplateData(c, "FlashSuccess", success)
		}
		if errMsg != "" {
			setTemplateData(c, "FlashError", errMsg)
		}

		c.Locals("flash_success", success)
		c.Locals("flash_error", errMsg)
		return c.Next()
	}
}

func setFlash(c *fiber.Ctx, key, message string) error {
	store, ok := c.Locals(contextStoreKey).(*session.Store)
	if !ok || store == nil {
		return errors.New("session store not available")
	}

	sess, err := store.Get(c)
	if err != nil {
		return err
	}

	sess.Set(key, message)
	return sess.Save()
}