// Package jwks fetches JSON Web Key Sets and validates RS256 JWTs against them.
//
// A Fetcher caches parsed RSA public keys in memory for 5 minutes per kid so
// repeated token validation does not re-hit the issuer's JWKS endpoint. It is
// safe for concurrent use.
package jwks

import (
	"context"
	"crypto/rsa"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"
	"net/http"
	"sync"
	"time"

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

// cacheTTL is how long a cached public key is considered fresh before the
// Fetcher re-fetches the JWKS document.
const cacheTTL = 5 * time.Minute

// Claims is the set of assertions carried by an Edu access token. The
// embedded RegisteredClaims provides exp, iss, aud and the other standard
// JWT claims.
type Claims struct {
	UserID    string `json:"user_id"`
	Role      string `json:"role"`
	DataScope string `json:"data_scope"`
	jwt.RegisteredClaims
}

// cachedKey holds a parsed public key together with the moment it expires.
type cachedKey struct {
	key       *rsa.PublicKey
	expiresAt time.Time
}

// Fetcher downloads and caches a JWKS document and validates tokens against
// the cached keys.
type Fetcher struct {
	jwksURL    string
	httpClient *http.Client
	cache      sync.Map
}

// NewFetcher returns a Fetcher that retrieves keys from jwksURL.
func NewFetcher(jwksURL string) *Fetcher {
	return &Fetcher{
		jwksURL:    jwksURL,
		httpClient: &http.Client{Timeout: 10 * time.Second},
	}
}

// GetPublicKey returns the RSA public key identified by kid. The key is served
// from the in-memory cache when fresh; otherwise the JWKS document is
// re-fetched and the cache repopulated.
func (f *Fetcher) GetPublicKey(kid string) (*rsa.PublicKey, error) {
	if v, ok := f.cache.Load(kid); ok {
		if ck, ok := v.(*cachedKey); ok && time.Now().Before(ck.expiresAt) {
			return ck.key, nil
		}
	}

	if err := f.refreshCache(); err != nil {
		return nil, err
	}

	v, ok := f.cache.Load(kid)
	if !ok {
		return nil, fmt.Errorf("jwks: kid %q not found in JWKS", kid)
	}
	ck, ok := v.(*cachedKey)
	if !ok {
		return nil, fmt.Errorf("jwks: cached entry for kid %q has unexpected type", kid)
	}
	return ck.key, nil
}

// refreshCache downloads the JWKS document and replaces every cached key.
func (f *Fetcher) refreshCache() error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, f.jwksURL, nil)
	if err != nil {
		return fmt.Errorf("jwks: building request: %w", err)
	}

	resp, err := f.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("jwks: fetching JWKS: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("jwks: unexpected status %d from %s", resp.StatusCode, f.jwksURL)
	}

	var body struct {
		Keys []struct {
			Kid string `json:"kid"`
			Kty string `json:"kty"`
			N   string `json:"n"`
			E   string `json:"e"`
		} `json:"keys"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
		return fmt.Errorf("jwks: decoding response: %w", err)
	}

	expiresAt := time.Now().Add(cacheTTL)
	for _, k := range body.Keys {
		if k.Kty != "RSA" {
			continue
		}
		pk, err := buildRSAPublicKey(k.N, k.E)
		if err != nil {
			return fmt.Errorf("jwks: parsing key %q: %w", k.Kid, err)
		}
		f.cache.Store(k.Kid, &cachedKey{key: pk, expiresAt: expiresAt})
	}
	return nil
}

// buildRSAPublicKey reconstructs an rsa.PublicKey from the base64url-encoded
// modulus (n) and exponent (e) fields of a JWK.
func buildRSAPublicKey(nStr, eStr string) (*rsa.PublicKey, error) {
	nBytes, err := base64.RawURLEncoding.DecodeString(nStr)
	if err != nil {
		return nil, fmt.Errorf("decoding modulus: %w", err)
	}
	eBytes, err := base64.RawURLEncoding.DecodeString(eStr)
	if err != nil {
		return nil, fmt.Errorf("decoding exponent: %w", err)
	}

	var e int
	for _, b := range eBytes {
		e = e<<8 + int(b)
	}
	if e == 0 {
		return nil, fmt.Errorf("empty exponent")
	}

	return &rsa.PublicKey{
		N: new(big.Int).SetBytes(nBytes),
		E: e,
	}, nil
}

// ValidateToken parses and validates an RS256 JWT using the cached JWKS keys.
// The kid in the token header selects the verification key. On success the
// populated Claims (including user_id, role, data_scope and the standard
// registered claims) are returned.
func (f *Fetcher) ValidateToken(token string) (*Claims, error) {
	parsed, err := jwt.ParseWithClaims(token, &Claims{}, func(t *jwt.Token) (any, error) {
		if _, ok := t.Method.(*jwt.SigningMethodRSA); !ok {
			return nil, fmt.Errorf("jwks: unexpected signing method %v", t.Header["alg"])
		}
		kid, ok := t.Header["kid"].(string)
		if !ok {
			return nil, fmt.Errorf("jwks: token header missing kid")
		}
		pk, err := f.GetPublicKey(kid)
		if err != nil {
			return nil, err
		}
		return pk, nil
	})
	if err != nil {
		return nil, fmt.Errorf("jwks: parsing token: %w", err)
	}

	claims, ok := parsed.Claims.(*Claims)
	if !ok || !parsed.Valid {
		return nil, fmt.Errorf("jwks: invalid token claims")
	}
	return claims, nil
}