package middleware

import (
	"net/http"
	"net/http/pprof"

	"github.com/labstack/echo/v4"
	"github.com/labstack/echo/v4/middleware"
)

type PprofConfig struct {
	Skipper middleware.Skipper
	Prefix  string
}

var DefaultPprofConfig = PprofConfig{
	Skipper: middleware.DefaultSkipper,
	Prefix:  "/debug/pprof",
}

func RegisterPprof(e *echo.Echo) {
	RegisterPprofWithConfig(e, DefaultPprofConfig)
}

func RegisterPprofWithConfig(e *echo.Echo, config PprofConfig, mw ...echo.MiddlewareFunc) {
	if config.Skipper == nil {
		config.Skipper = DefaultPprofConfig.Skipper
	}
	if config.Prefix == "" {
		config.Prefix = DefaultPprofConfig.Prefix
	}

	g := e.Group(config.Prefix)
	{
		g.GET("/", pprofAdapter(pprof.Index), mw...)
		g.GET("/allocs", pprofAdapter(pprof.Handler("allocs").ServeHTTP), mw...)
		g.GET("/block", pprofAdapter(pprof.Handler("block").ServeHTTP), mw...)
		g.GET("/cmdline", pprofAdapter(pprof.Cmdline), mw...)
		g.GET("/goroutine", pprofAdapter(pprof.Handler("goroutine").ServeHTTP), mw...)
		g.GET("/heap", pprofAdapter(pprof.Handler("heap").ServeHTTP), mw...)
		g.GET("/mutex", pprofAdapter(pprof.Handler("mutex").ServeHTTP), mw...)
		g.GET("/profile", pprofAdapter(pprof.Profile), mw...)
		g.POST("/symbol", pprofAdapter(pprof.Symbol), mw...)
		g.GET("/symbol", pprofAdapter(pprof.Symbol), mw...)
		g.GET("/threadcreate", pprofAdapter(pprof.Handler("threadcreate").ServeHTTP), mw...)
		g.GET("/trace", pprofAdapter(pprof.Trace), mw...)
	}
}

func pprofAdapter(h http.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		h.ServeHTTP(c.Response().Writer, c.Request())
		return nil
	}
}