package mymiddleware

import (
	"encoding/base64"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

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

const (
	// DefaultTokenLookup defines the default location to extract JWT tokens from
	DefaultTokenLookup = "header:Authorization"
	// DefaultAuthScheme defines the default authentication scheme
	DefaultAuthScheme = "Bearer"
	// DefaultContextKey defines the default key for storing full JWT claims in Echo context
	DefaultContextKey = "user"
	// DefaultEmailContextKey defines the key for storing extracted email in Echo context
	DefaultEmailContextKey = "user_email"
	// DefaultEmailKey defines the default key for extracting email from JWT claims
	DefaultEmailKey = "email"
)

// JWTConfig holds the configuration for JWT middleware
type JWTConfig struct {
	// TokenLookup defines where to extract the token from (default: "header:Authorization")
	TokenLookup string
	// AuthScheme defines the auth scheme (default: "Bearer")
	AuthScheme string
	// ContextKey is the key used to store user data in context (default: "user")
	ContextKey string
	// UserEmail is the key to extract email from JWT claims (default: "email")
	UserEmail string
	// Skipper defines a function to skip middleware
	Skipper func(echo.Context) bool
}

// DefaultJWTConfig returns a default JWT configuration
func DefaultJWTConfig() JWTConfig {
	return JWTConfig{
		TokenLookup: DefaultTokenLookup,
		AuthScheme:  DefaultAuthScheme,
		ContextKey:  DefaultContextKey,
		UserEmail:   DefaultEmailKey,
		Skipper:     nil,
	}
}

// JWTWithConfig returns a JWT middleware with custom configuration
func JWTWithConfig(config JWTConfig) echo.MiddlewareFunc {
	// Set defaults
	if config.TokenLookup == "" {
		config.TokenLookup = DefaultTokenLookup
	}
	if config.AuthScheme == "" {
		config.AuthScheme = DefaultAuthScheme
	}
	if config.ContextKey == "" {
		config.ContextKey = DefaultContextKey
	}
	if config.UserEmail == "" {
		config.UserEmail = DefaultEmailKey
	}

	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Skip middleware if skipper function returns true
			if config.Skipper != nil && config.Skipper(c) {
				return next(c)
			}

			// Extract token from request
			token, err := extractToken(c, config)
			if err != nil {
				return echo.NewHTTPError(http.StatusUnauthorized, "missing or invalid token")
			}

			// Parse JWT payload without signature verification
			claims, err := parseJWTPayloadUnsafe(token)
			if err != nil {
				return echo.NewHTTPError(http.StatusUnauthorized, "invalid token: "+err.Error())
			}

			// Extract email from claims
			email, exists := claims[config.UserEmail]
			if !exists {
				return echo.NewHTTPError(http.StatusUnauthorized, "email not found in token")
			}

			emailStr, ok := email.(string)
			if !ok {
				return echo.NewHTTPError(http.StatusUnauthorized, "invalid email format in token")
			}

			// INFO: temporary restriction for email
			if !strings.HasSuffix(emailStr, "@pingcap.com") {
				return echo.NewHTTPError(http.StatusUnauthorized, "invalid email format in token")
			}

			// Set user data in context
			c.Set(config.ContextKey, claims)
			c.Set(DefaultEmailContextKey, emailStr)

			return next(c)
		}
	}
}

// JWT returns a JWT middleware with default configuration
func JWT() echo.MiddlewareFunc {
	config := DefaultJWTConfig()
	return JWTWithConfig(config)
}

// extractToken extracts JWT token from the request based on configuration
func extractToken(c echo.Context, config JWTConfig) (string, error) {
	// Currently only supports header extraction
	if !strings.HasPrefix(config.TokenLookup, "header:") {
		return "", fmt.Errorf("unsupported token lookup method")
	}

	header := strings.TrimPrefix(config.TokenLookup, "header:")
	auth := c.Request().Header.Get(header)
	if auth == "" {
		return "", fmt.Errorf("missing authorization header")
	}

	// Extract token from "Bearer <token>" format
	authParts := strings.SplitN(auth, " ", 2)
	if len(authParts) != 2 || authParts[0] != config.AuthScheme {
		return "", fmt.Errorf("invalid authorization header format")
	}

	return authParts[1], nil
}

// GetUserEmailFromContext extracts the email from Echo context
func GetUserEmailFromContext(c echo.Context) (string, bool) {
	email, ok := c.Get(DefaultEmailContextKey).(string)
	return email, ok
}

// parseJWTPayloadUnsafe parses JWT payload without signature verification
// WARNING: This is unsafe and should only be used when you trust the token source
func parseJWTPayloadUnsafe(tokenString string) (map[string]interface{}, error) {
	// Split token into parts
	parts := strings.Split(tokenString, ".")
	if len(parts) != 3 {
		return nil, fmt.Errorf("invalid token format")
	}

	// Decode payload (second part)
	payload := parts[1]

	// Add padding if necessary for base64 decoding
	if l := len(payload) % 4; l > 0 {
		payload += strings.Repeat("=", 4-l)
	}

	// Base64 decode the payload
	payloadBytes, err := base64.URLEncoding.DecodeString(payload)
	if err != nil {
		return nil, fmt.Errorf("failed to decode payload: %w", err)
	}

	// Unmarshal JSON payload
	var claims map[string]interface{}
	if err := json.Unmarshal(payloadBytes, &claims); err != nil {
		return nil, fmt.Errorf("failed to unmarshal payload: %w", err)
	}

	return claims, nil
}

// GetUserClaimsFromContext extracts the full JWT claims from Echo context
func GetUserClaimsFromContext(c echo.Context) (map[string]interface{}, bool) {
	claims, ok := c.Get(DefaultContextKey).(map[string]interface{})
	return claims, ok
}