package middleware

import (
	"context"
	"fmt"
	"net/http"
	"strings"

	"github.com/golang-jwt/jwt/v5"
)

type contextKey string

const (
	KeySubject contextKey = "subject"
	KeyScopes  contextKey = "scopes"
)

// Auth 验证 Bearer Token（API Key 或 JWT）。
type Auth struct {
	apiKeys   map[string]string // key → description
	jwtSecret []byte
}

func NewAuth(apiKeys map[string]string, jwtSecret string) *Auth {
	return &Auth{
		apiKeys:   apiKeys,
		jwtSecret: []byte(jwtSecret),
	}
}

// Middleware 从 HTTP Header 中提取并验证 Token。
func (a *Auth) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		token := extractBearer(r)
		if token == "" {
			http.Error(w, "missing Bearer token", http.StatusUnauthorized)
			return
		}

		// 尝试 API Key 验证
		if desc, ok := a.apiKeys[token]; ok {
			ctx := context.WithValue(r.Context(), KeySubject, fmt.Sprintf("apikey:%s", desc))
			ctx = context.WithValue(ctx, KeyScopes, []string{"*"})
			next.ServeHTTP(w, r.WithContext(ctx))
			return
		}

		// 尝试 JWT 验证
		if a.jwtSecret != nil {
			claims, err := a.parseJWT(token)
			if err == nil {
				ctx := context.WithValue(r.Context(), KeySubject, claims.Subject)
				ctx = context.WithValue(ctx, KeyScopes, claims.Scopes)
				next.ServeHTTP(w, r.WithContext(ctx))
				return
			}
		}

		http.Error(w, "invalid token", http.StatusUnauthorized)
	})
}

type customClaims struct {
	jwt.RegisteredClaims
	Scopes []string `json:"scopes"`
}

func (a *Auth) parseJWT(tokenStr string) (*customClaims, error) {
	token, err := jwt.ParseWithClaims(tokenStr, &customClaims{},
		func(t *jwt.Token) (any, error) {
			return a.jwtSecret, nil
		},
	)
	if err != nil {
		return nil, err
	}
	claims, ok := token.Claims.(*customClaims)
	if !ok || !token.Valid {
		return nil, fmt.Errorf("invalid claims")
	}
	return claims, nil
}

func extractBearer(r *http.Request) string {
	auth := r.Header.Get("Authorization")
	if auth == "" {
		return ""
	}
	if !strings.HasPrefix(auth, "Bearer ") {
		return ""
	}
	return strings.TrimPrefix(auth, "Bearer ")
}