package trace

import (
	"context"
	"loveuer/utodo/pkg/logger"
	"strconv"
	"time"

	"github.com/gofiber/fiber/v3"
	"github.com/google/uuid"
)

type Option func(*option)

type option struct {
	traceKey string
	traceFn  func(c fiber.Ctx) string
}

func WithTraceKey(key string) Option {
	return func(o *option) {
		if key != "" {
			o.traceKey = key
		}
	}
}

func WithTraceFn(fn func(c fiber.Ctx) string) Option {
	return func(o *option) {
		if fn != nil {
			o.traceFn = fn
		}
	}
}

func New(opts ...Option) fiber.Handler {
	opt := &option{
		traceKey: "X-Trace-Id",
		traceFn: func(c fiber.Ctx) string {
			uid, err := uuid.NewV7()
			if err != nil {
				return strconv.FormatInt(time.Now().UnixNano(), 10)
			}

			return uid.String()
		},
	}

	for _, o := range opts {
		o(opt)
	}

	return func(c fiber.Ctx) error {
		traceId := c.Get(opt.traceKey)
		if traceId == "" {
			traceId = opt.traceFn(c)
			c.Request().Header.Set(opt.traceKey, traceId)
		}

		c.SetContext(context.WithValue(c.Context(), logger.CtxKey, traceId))
		c.Set(opt.traceKey, traceId)

		return c.Next()
	}
}