package main

import (
	"encoding/json"
	"net/http"
	"os"
	"strings"

	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt"
	"github.com/pingcap/log"
	"go.uber.org/zap"
)

const (
	envJwtAuthEnabled = "O11Y_JWT_AUTH_ENABLED"
	envJwtSecret      = "JWT_SECRET"

	headerO11yToken = "O11ytoken"
	// #nosec G101 -- header name for o11y JWT token, not a credential
	headerXO11yToken = "X-O11y-Token"
	queryO11yToken   = "o11ytoken"

	contextKeyUser = "github_username"
)

type authServerVerifyResp struct {
	GithubUsername string `json:"github_username"`
}

type tokenVerifyError struct {
	statusCode int
	message    string
}

func (e *tokenVerifyError) Error() string {
	return e.message
}

// o11yJwtAuthMiddleware returns a Gin middleware. When O11Y_JWT_AUTH_ENABLED is not "true",
// it does nothing. When enabled, it requires a valid o11ytoken and validates it locally
// via JWT_SECRET. The token must include github_username claim.
func o11yJwtAuthMiddleware() gin.HandlerFunc {
	enabled := strings.ToLower(strings.TrimSpace(os.Getenv(envJwtAuthEnabled))) == "true"
	jwtSecret := strings.TrimSpace(os.Getenv(envJwtSecret))

	if enabled && jwtSecret == "" {
		log.Warn("o11y jwt auth: JWT_SECRET is empty while auth is enabled")
	}

	return func(c *gin.Context) {
		if !enabled {
			c.Next()
			return
		}

		tokenStr := c.GetHeader(headerO11yToken)
		if tokenStr == "" {
			tokenStr = c.GetHeader(headerXO11yToken)
		}
		if tokenStr == "" {
			tokenStr = c.Query(queryO11yToken)
		}
		if tokenStr == "" {
			log.Warn("o11y jwt auth: missing o11ytoken")
			c.JSON(http.StatusUnauthorized, gin.H{"error": "missing o11ytoken"})
			c.Abort()
			return
		}

		verifyResp, err := verifyTokenLocally(tokenStr, jwtSecret)
		if err != nil {
			log.Warn("o11y jwt auth: local verify failed", zap.Error(err))
			c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid token"})
			c.Abort()
			return
		}

		githubUsername := strings.TrimSpace(verifyResp.GithubUsername)
		if githubUsername == "" {
			log.Warn("o11y jwt auth: missing github_username in verify response")
			c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid token"})
			c.Abort()
			return
		}

		c.Set(contextKeyUser, githubUsername)

		// Log request params and username for statistics.
		params := map[string]interface{}{
			"method":          c.Request.Method,
			"path":            c.Request.URL.Path,
			"raw_query":       c.Request.URL.RawQuery,
			"github_username": githubUsername,
		}
		if c.Request.URL.RawQuery != "" {
			params["query"] = c.Request.URL.Query()
		}
		paramsJSON, _ := json.Marshal(params)
		log.Info("o11y jwt auth: request", zap.String("params", string(paramsJSON)))

		c.Next()
	}
}

func verifyTokenLocally(tokenStr, jwtSecret string) (*authServerVerifyResp, error) {
	if strings.TrimSpace(jwtSecret) == "" {
		return nil, &tokenVerifyError{statusCode: http.StatusUnauthorized, message: "invalid token"}
	}
	token, err := jwt.Parse(tokenStr, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
			return nil, &tokenVerifyError{statusCode: http.StatusUnauthorized, message: "invalid token"}
		}
		return []byte(jwtSecret), nil
	})
	if err != nil {
		return nil, err
	}
	if !token.Valid {
		return nil, &tokenVerifyError{statusCode: http.StatusUnauthorized, message: "invalid token"}
	}
	claims, ok := token.Claims.(jwt.MapClaims)
	if !ok {
		return nil, &tokenVerifyError{statusCode: http.StatusUnauthorized, message: "invalid token"}
	}
	githubUsername, ok := claims[contextKeyUser].(string)
	if !ok || strings.TrimSpace(githubUsername) == "" {
		return nil, &tokenVerifyError{statusCode: http.StatusUnauthorized, message: "invalid token"}
	}
	return &authServerVerifyResp{GithubUsername: strings.TrimSpace(githubUsername)}, nil
}