package middleware

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

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

type contextKey string
const authKey contextKey = "auth"

type AuthClaims struct {
    UserID  string
    OrgID   string
    OrgRole string
}

func Auth(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) {
            token := bearerToken(r)
            if token == "" {
                http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized)
                return
            }

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

            claims := AuthClaims{
                UserID:  payload.Sub,
                OrgID:   payload.OrgID,
                OrgRole: payload.OrgRole,
            }
            ctx := context.WithValue(r.Context(), authKey, claims)
            next.ServeHTTP(w, r.WithContext(ctx))
        })
    }
}

func GetAuth(r *http.Request) (AuthClaims, bool) {
    claims, ok := r.Context().Value(authKey).(AuthClaims)
    return claims, ok
}

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

mux := http.NewServeMux()
auth := middleware.Auth(hj)

mux.Handle("GET /api/me", auth(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    claims, _ := middleware.GetAuth(r)
    json.NewEncoder(w).Encode(map[string]string{"userId": claims.UserID})
})))

import "github.com/go-chi/chi/v5"

r := chi.NewRouter()
r.Use(middleware.Auth(hj)) // Apply globally

r.Get("/api/me", func(w http.ResponseWriter, r *http.Request) {
    claims, _ := middleware.GetAuth(r)
    json.NewEncoder(w).Encode(map[string]string{"userId": claims.UserID})
})

// Group with auth
r.Group(func(r chi.Router) {
    r.Use(middleware.Auth(hj))
    r.Get("/api/protected", protectedHandler)
    r.Post("/api/data", createDataHandler)
})

import "github.com/gin-gonic/gin"

func HJAuth(hj *hellojohn.Client) gin.HandlerFunc {
    return func(c *gin.Context) {
        token := strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer ")
        if token == "" {
            c.JSON(http.StatusUnauthorized, gin.H{"error": "missing token"})
            c.Abort()
            return
        }

        payload, err := hj.VerifyToken(c.Request.Context(), token)
        if err != nil {
            c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid token"})
            c.Abort()
            return
        }

        c.Set("userId", payload.Sub)
        c.Set("orgId", payload.OrgID)
        c.Next()
    }
}

// Usage
r := gin.Default()
protected := r.Group("/api", HJAuth(hj))
protected.GET("/me", func(c *gin.Context) {
    userID := c.GetString("userId")
    c.JSON(200, gin.H{"userId": userID})
})

func RequireRole(role string) func(http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            claims, ok := middleware.GetAuth(r)
            if !ok {
                http.Error(w, `{"error":"unauthorized"}`, http.StatusUnauthorized)
                return
            }
            if claims.OrgRole != role && claims.OrgRole != "owner" {
                http.Error(w, `{"error":"forbidden"}`, http.StatusForbidden)
                return
            }
            next.ServeHTTP(w, r)
        })
    }
}

// Usage: require admin role
mux.Handle("DELETE /api/users/{id}",
    auth(RequireRole("admin")(http.HandlerFunc(deleteUserHandler))))