package httpserver

import (
	"github.com/julienschmidt/httprouter"
	"net/http"
	"net/http/pprof"
)

var pprofFuncs = map[string]http.HandlerFunc{
	"/debug/pprof/":          pprof.Index,
	"/debug/pprof/goroutine": pprof.Handler("goroutine").ServeHTTP,
	"/debug/pprof/heap":      pprof.Handler("heap").ServeHTTP,
	"/debug/pprof/allocs":    pprof.Handler("allocs").ServeHTTP,
	"/debug/pprof/cmdline":   pprof.Cmdline,
	"/debug/pprof/profile":   pprof.Profile,
	"/debug/pprof/symbol":    pprof.Symbol,
	"/debug/pprof/trace":     pprof.Trace,
}

//PProfRegRouter pprof-注册路由
//查看内存占用分析： go tool pprof -alloc_space "127.0.0.1:8888/debug/pprof/heap?seconds=30"
func PProfRegRouter(ro interface{}) {
	//默认不传则直接注册http
	if ro == nil {
		for path, f := range pprofFuncs {
			http.HandleFunc(path, f)
		}
		return
	}
	//支持本地 HttpServer 进行pprof处理
	if r, ok := ro.(HttpServer); ok {
		for path, f := range pprofFuncs {
			r.RouteHandler(http.MethodGet, path, f, nil)
		}
		return
	}
	//支持 http.Server 的处理
	if s, ok := ro.(*http.Server); ok {
		s.Handler = PProfHandler(s.Handler.ServeHTTP)
		return
	}
	//支持 httprouter.Router 的处理
	if r, ok := ro.(*httprouter.Router); ok {
		for path, f := range pprofFuncs {
			r.HandlerFunc(http.MethodGet, path, f)
		}
		return
	}
}

func PProfHandler(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		//针对pprof路径进行处理
		if f, ok := pprofFuncs[r.URL.Path]; ok {
			f(w, r)
			return
		}
		if next != nil {
			next(w, r)
		}
	}
}