package cleanup

import (
	"context"
	"log/slog"
	"net/http"
	"reflect"

	"git.jeffthecoder.xyz/public/lazyhandler/middleware"
)

type ctxKey int

const (
	cleanupCtxKey ctxKey = iota
)

type CleanupContext struct {
	funcs []Cleanup
}

type Cleanup interface {
	Name() string
	Cleanup()
}

type CleanupFunc func()

func (fn CleanupFunc) Name() string {
	return reflect.TypeOf(fn).String()
}

func (fn CleanupFunc) Cleanup() {
	defer func() {
		if v := recover(); v != nil {
			slog.With("v", v).Warn("cleanup panicked")
		}
	}()

	fn()
}

func Collect() middleware.Middleware {
	return middleware.WrapFunc(func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ctx := &CleanupContext{}
			next.ServeHTTP(w, r.WithContext(context.WithValue(r.Context(), cleanupCtxKey, ctx)))
			for _, cleanupFunc := range ctx.funcs {
				defer cleanupFunc.Cleanup()
			}
		})
	})
}

// Register adds a Cleanup function to the CleanupContext in the provided context.
// If the CleanupContext is not found in the context, the Cleanup function is not registered.
func Register(ctx context.Context, c Cleanup) {
	if ctx, ok := ctx.Value(cleanupCtxKey).(*CleanupContext); ok {
		ctx.funcs = append(ctx.funcs, c)
	}
}