package middleware

import (
	"bytes"
	"io"
	"net/http"
	"strings"

	"github.com/zero-ppanel/zero-ppanel/apps/admin/internal/config"
	"github.com/zero-ppanel/zero-ppanel/pkg/signature"
	"github.com/zero-ppanel/zero-ppanel/pkg/xerr"
	"github.com/zeromicro/go-zero/rest/httpx"
)

type SignatureMiddleware struct {
	conf      config.Config
	validator *signature.Validator
}

func NewSignatureMiddleware(c config.Config, store signature.NonceStore) *SignatureMiddleware {
	return &SignatureMiddleware{
		conf:      c,
		validator: signature.NewValidator(c.AppSignature, store),
	}
}

func (m *SignatureMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		appId := r.Header.Get("X-App-Id")
		if appId == "" {
			next(w, r)
			return
		}

		for _, prefix := range m.conf.AppSignature.SkipPrefixes {
			if strings.HasPrefix(r.URL.Path, prefix) {
				next(w, r)
				return
			}
		}

		timestamp := r.Header.Get("X-Timestamp")
		nonce := r.Header.Get("X-Nonce")
		sig := r.Header.Get("X-Signature")

		if timestamp == "" || nonce == "" || sig == "" {
			httpx.WriteJson(w, http.StatusUnauthorized, buildErrResp(xerr.SignatureMissing))
			return
		}

		var bodyBytes []byte
		if r.Body != nil {
			bodyBytes, _ = io.ReadAll(r.Body)
			r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
		}

		sts := signature.BuildStringToSign(r.Method, r.URL.Path, r.URL.RawQuery, bodyBytes, appId, timestamp, nonce)

		if err := m.validator.Validate(r.Context(), appId, timestamp, nonce, sig, sts); err != nil {
			code := mapSignatureErr(err)
			httpx.WriteJson(w, http.StatusUnauthorized, buildErrResp(code))
			return
		}

		next(w, r)
	}
}

func mapSignatureErr(err error) int {
	switch err {
	case signature.ErrSignatureMissing:
		return xerr.SignatureMissing
	case signature.ErrSignatureExpired:
		return xerr.SignatureExpired
	case signature.ErrSignatureReplay:
		return xerr.SignatureReplay
	default:
		return xerr.SignatureInvalid
	}
}

func buildErrResp(code int) map[string]interface{} {
	return map[string]interface{}{
		"code": code,
		"msg":  xerr.MapErrMsg(code),
		"data": nil,
	}
}