package middleware

import (
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"
	"net/http"
	"os"
	"strings"
	"sync"

	"github.com/gofiber/fiber/v2"
	"github.com/golang-jwt/jwt/v5"
)

var (
	jwksOnce  sync.Once
	jwksCache map[string]*rsa.PublicKey
)

func loadJWKS() map[string]*rsa.PublicKey {
	jwksOnce.Do(func() {
		region := os.Getenv("AWS_REGION")
		poolID := os.Getenv("COGNITO_USER_POOL_ID")
		url := fmt.Sprintf("https://cognito-idp.%s.amazonaws.com/%s/.well-known/jwks.json", region, poolID)
		resp, err := http.Get(url) //nolint:noctx
		if err != nil {
			return
		}
		defer resp.Body.Close()
		var body struct {
			Keys []struct {
				Kid string `json:"kid"`
				N   string `json:"n"`
				E   string `json:"e"`
			} `json:"keys"`
		}
		if json.NewDecoder(resp.Body).Decode(&body) != nil {
			return
		}
		jwksCache = make(map[string]*rsa.PublicKey, len(body.Keys))
		for _, k := range body.Keys {
			nBytes, _ := base64.RawURLEncoding.DecodeString(k.N)
			eBytes, _ := base64.RawURLEncoding.DecodeString(k.E)
			jwksCache[k.Kid] = &rsa.PublicKey{
				N: new(big.Int).SetBytes(nBytes),
				E: int(new(big.Int).SetBytes(eBytes).Int64()),
			}
		}
	})
	return jwksCache
}

func RequireAuth(c *fiber.Ctx) error {
	header := c.Get("Authorization")
	if !strings.HasPrefix(header, "Bearer ") {
		return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "unauthorized"})
	}
	tokenStr := strings.TrimPrefix(header, "Bearer ")

	keys := loadJWKS()
	token, err := jwt.Parse(tokenStr, func(t *jwt.Token) (any, error) {
		if _, ok := t.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
		}
		kid, _ := t.Header["kid"].(string)
		key, ok := keys[kid]
		if !ok {
			return nil, fmt.Errorf("unknown kid: %s", kid)
		}
		return key, nil
	})
	if err != nil || !token.Valid {
		return c.Status(fiber.StatusUnauthorized).JSON(fiber.Map{"error": "invalid token"})
	}

	c.Locals("claims", token.Claims)
	return c.Next()
}
