package main

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

const (
	defaultJWKSCacheTTL        = time.Hour
	defaultJWKSRefreshCooldown = time.Minute
)

type jwksDocument struct {
	Keys []jsonWebKey `json:"keys"`
}

type jsonWebKey struct {
	Algorithm string `json:"alg"`
	Exponent  string `json:"e"`
	KeyID     string `json:"kid"`
	KeyType   string `json:"kty"`
	Modulus   string `json:"n"`
	Use       string `json:"use"`
}

type jwksCache struct {
	url        string
	httpClient *http.Client
	ttl        time.Duration

	mu        sync.Mutex
	keys      map[string]*rsa.PublicKey
	expiresAt time.Time

	now                func() time.Time
	nextRefreshAllowed time.Time
	refreshCooldown    time.Duration
}

func newJWKSCache(url string, httpClient *http.Client, ttl time.Duration) *jwksCache {
	if ttl <= 0 {
		ttl = defaultJWKSCacheTTL
	}
	return &jwksCache{
		url:             url,
		httpClient:      httpClient,
		ttl:             ttl,
		now:             time.Now,
		refreshCooldown: defaultJWKSRefreshCooldown,
	}
}

func (c *jwksCache) key(ctx context.Context, keyID string) (*rsa.PublicKey, error) {
	keyID = strings.TrimSpace(keyID)
	if keyID == "" {
		return nil, fmt.Errorf("FIT header does not contain kid")
	}

	c.mu.Lock()
	defer c.mu.Unlock()

	now := c.now()
	if key := c.keys[keyID]; key != nil && now.Before(c.expiresAt) {
		return key, nil
	}
	if now.Before(c.nextRefreshAllowed) {
		return nil, fmt.Errorf("JWKS does not contain key %q", keyID)
	}

	// Limit both refresh-on-miss and failed refresh attempts. Without this
	// guard, requests carrying random kid values can force one outbound JWKS
	// request each even while the existing cache is still fresh.
	c.nextRefreshAllowed = now.Add(c.refreshCooldown)
	if err := c.refreshLocked(ctx); err != nil {
		return nil, err
	}
	key := c.keys[keyID]
	if key == nil {
		return nil, fmt.Errorf("JWKS does not contain key %q", keyID)
	}
	return key, nil
}

func (c *jwksCache) refreshLocked(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.url, nil)
	if err != nil {
		return fmt.Errorf("create JWKS request: %w", err)
	}
	req.Header.Set("Accept", "application/json")

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("fetch JWKS: %w", err)
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("fetch JWKS: unexpected status %d", resp.StatusCode)
	}

	var document jwksDocument
	if err := json.NewDecoder(resp.Body).Decode(&document); err != nil {
		return fmt.Errorf("decode JWKS: %w", err)
	}

	keys := make(map[string]*rsa.PublicKey, len(document.Keys))
	for _, item := range document.Keys {
		if item.KeyType != "RSA" || (item.Use != "" && item.Use != "sig") || (item.Algorithm != "" && item.Algorithm != "RS256") {
			continue
		}
		key, err := rsaPublicKey(item)
		if err != nil {
			return fmt.Errorf("decode JWKS key %q: %w", item.KeyID, err)
		}
		keys[item.KeyID] = key
	}
	if len(keys) == 0 {
		return fmt.Errorf("JWKS does not contain an RS256 signing key")
	}

	c.keys = keys
	c.expiresAt = c.now().Add(c.ttl)
	return nil
}

func rsaPublicKey(key jsonWebKey) (*rsa.PublicKey, error) {
	modulus, err := decodeBase64URL(key.Modulus)
	if err != nil {
		return nil, fmt.Errorf("decode modulus: %w", err)
	}
	exponentBytes, err := decodeBase64URL(key.Exponent)
	if err != nil {
		return nil, fmt.Errorf("decode exponent: %w", err)
	}

	exponent := 0
	for _, value := range exponentBytes {
		exponent = exponent<<8 + int(value)
	}
	if len(modulus) == 0 || exponent < 2 {
		return nil, fmt.Errorf("invalid RSA public key")
	}

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

func decodeBase64URL(value string) ([]byte, error) {
	decoded, err := base64.RawURLEncoding.DecodeString(value)
	if err == nil {
		return decoded, nil
	}
	return base64.URLEncoding.DecodeString(value)
}