package middleware

import (
	"context"
	"fmt"
	"github.com/labstack/echo/v4"
	"github.com/rs/xid"
	"net/http"
	"sorbet/internal/util"
	"sorbet/pkg/log"
	"time"
)

var color = log.NewColorer()

func requestId(req *http.Request, res *echo.Response) string {
	id := req.Header.Get(echo.HeaderXRequestID)
	if id == "" {
		id = xid.New().String()
		res.Header().Set(echo.HeaderXRequestID, id)
	}
	return id
}

// Logger 该日志中间件会自动获取或设置 RequestID
func Logger(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) (err error) {
		req := c.Request()
		res := c.Response()
		start := time.Now()
		id := requestId(req, res)
		l := log.With(log.String("id", id))
		l.Info(
			"Started %s %s for %s",
			req.Method, req.RequestURI, c.RealIP(),
			log.RawTime(start),
		)
		c.SetRequest(c.Request().WithContext(context.WithValue(c.Request().Context(), "logger", l)))
		c.SetLogger(util.NewCustomLogger(l))
		if err = next(c); err != nil {
			c.Error(err)
		}
		stop := time.Now()
		content := fmt.Sprintf(
			"Completed %s %s %v %s in %v",
			req.Method, req.RequestURI, res.Status,
			http.StatusText(res.Status), stop.Sub(start),
		)
		if res.Status >= 500 {
			content = color.Cyan(content)
		} else if res.Status >= 400 {
			content = color.Red(content)
		} else if res.Status >= 300 {
			if res.Status == 304 {
				content = color.Yellow(content)
			} else {
				content = color.White(content)
			}
		} else if res.Status >= 200 {
			content = color.Green(content)
		}
		l.Info(content, log.RawTime(stop))
		return
	}
}