import "github.com/hellojohn/hellojohn-go"

hj := hellojohn.New(hellojohn.Config{
    TenantID:  os.Getenv("HELLOJOHN_TENANT_ID"),
    SecretKey: os.Getenv("HELLOJOHN_SECRET_KEY"),
})

payload, err := hj.VerifyToken(ctx, token)
if err != nil {
    // Invalid, expired, or malformed token
    http.Error(w, "Unauthorized", http.StatusUnauthorized)
    return
}

fmt.Println(payload.Sub) // User ID (sub claim)

type TokenPayload struct {
    Sub     string    `json:"sub"`      // User ID
    Sid     string    `json:"sid"`      // Session ID
    Iss     string    `json:"iss"`      // Issuer
    Aud     string    `json:"aud"`      // Audience (tenant ID)
    Iat     int64     `json:"iat"`      // Issued at
    Exp     int64     `json:"exp"`      // Expiry
    OrgID   string    `json:"org_id"`   // Active org ID (if any)
    OrgRole string    `json:"org_role"` // Role in active org
}

func AuthMiddleware(hj *hellojohn.Client) func(http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            authHeader := r.Header.Get("Authorization")
            if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") {
                http.Error(w, `{"error":"missing token"}`, http.StatusUnauthorized)
                return
            }

            token := strings.TrimPrefix(authHeader, "Bearer ")
            payload, err := hj.VerifyToken(r.Context(), token)
            if err != nil {
                http.Error(w, `{"error":"invalid token"}`, http.StatusUnauthorized)
                return
            }

            // Store user ID in request context
            ctx := context.WithValue(r.Context(), userIDKey, payload.Sub)
            next.ServeHTTP(w, r.WithContext(ctx))
        })
    }
}

// Retrieve from context
func getUserID(r *http.Request) string {
    return r.Context().Value(userIDKey).(string)
}

mux := http.NewServeMux()

// Public
mux.HandleFunc("GET /api/health", func(w http.ResponseWriter, r *http.Request) {
    w.Write([]byte(`{"ok":true}`))
})

// Protected — wrap with middleware
protected := AuthMiddleware(hj)(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    userID := getUserID(r)
    json.NewEncoder(w).Encode(map[string]string{"userId": userID})
}))

mux.Handle("GET /api/me", protected)

import (
    "github.com/go-jose/go-jose/v4"
    "github.com/go-jose/go-jose/v4/jwt"
)

// Fetch JWKS once and cache
jwksURL := fmt.Sprintf(
    "https://api.hellojohn.dev/v1/jwks/%s",
    os.Getenv("HELLOJOHN_TENANT_ID"),
)

// Parse and verify
keySet := // fetch and parse JWKS...

tok, err := jwt.ParseSigned(tokenStr, []jose.SignatureAlgorithm{jose.EdDSA})
if err != nil {
    // invalid
}

var claims jwt.Claims
if err := tok.Claims(keySet, &claims); err != nil {
    // signature invalid
}

if err := claims.ValidateWithLeeway(jwt.Expected{
    Issuer:   "https://api.hellojohn.dev",
    Audience: jwt.Audience{os.Getenv("HELLOJOHN_TENANT_ID")},
    Time:     time.Now(),
}, time.Second*30); err != nil {
    // expired or wrong issuer/audience
}