package cognitoauth

import (
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"time"

	"github.com/lestrrat-go/jwx/v2/jwk"
)

// Time constants for JWKS operations
const (
	// JWKS_CACHE_TTL is the time-to-live duration for the JWKS cache (24 hours)
	JWKS_CACHE_TTL = 24 * time.Hour

	// JWKS_CLIENT_TIMEOUT is the timeout duration for HTTP requests to fetch JWKS (10 seconds)
	JWKS_CLIENT_TIMEOUT = 10 * time.Second
)

// GetJWKS retrieves a JSON Web Key Set (JWKS) from the specified URL, using a
// thread-safe caching mechanism to minimize external requests.
//
// It implements a double-checked locking pattern to efficiently handle concurrent
// access to the cached JWKS. It first attempts to retrieve a valid cached key set
// using a read lock. If the cache is empty or expired, it acquires a write lock
// to update the cache, but does a second check to avoid redundant fetches in case
// another goroutine has already updated the cache.
//
// The retrieved JWKS is cached for JWKS_CACHE_TTL duration (24 hours by default).
// This cache duration can be adjusted based on the security requirements and the
// key rotation policy of the identity provider.
//
// Parameters:
//   - jwksURL: The URL from which to fetch the JWKS (typically provided by the Cognito or OAuth provider)
//   - logger: A structured logger for recording debug and error information during the JWKS retrieval process
//
// Returns:
//   - jwk.Set: The parsed JSON Web Key Set containing the public keys used for token verification
//   - error: An error if the JWKS cannot be fetched or parsed, nil otherwise
//
// This function is primarily used during token validation to get the key set
// needed for signature verification of JWTs issued by the identity provider.
func GetJWKS(jwksURL string, logger *slog.Logger) (jwk.Set, error) {
	// Try to use cached JWKS first (read lock)
	jwksCache.mutex.RLock()
	if jwksCache.KeySet != nil && time.Now().Before(jwksCache.ExpiresAt) {
		defer jwksCache.mutex.RUnlock()
		logger.Debug("Using cached JWKS")
		return jwksCache.KeySet, nil
	}
	jwksCache.mutex.RUnlock()

	// Need to fetch new JWKS (write lock)
	jwksCache.mutex.Lock()
	defer jwksCache.mutex.Unlock()

	// Double-check expiration after acquiring the write lock
	if jwksCache.KeySet != nil && time.Now().Before(jwksCache.ExpiresAt) {
		return jwksCache.KeySet, nil
	}

	logger.Info("Fetching new JWKS", "url", jwksURL)
	jwksJSON, err := fetchJWKS(jwksURL)
	if err != nil {
		return nil, fmt.Errorf("failed to fetch JWKS: %w", err)
	}

	keySet, err := jwk.ParseString(jwksJSON)
	if err != nil {
		return nil, fmt.Errorf("failed to parse JWKS: %w", err)
	}

	jwksCache.KeySet = keySet
	// Cache for 24 hours - you can adjust this based on your needs
	jwksCache.ExpiresAt = time.Now().Add(JWKS_CACHE_TTL)

	return keySet, nil
}

// fetchJWKS makes an HTTP request to retrieve the JSON Web Key Set (JWKS) from
// the specified URL, typically from a Cognito or other OAuth identity provider.
//
// This function handles the actual HTTP communication with the JWKS endpoint,
// using a timeout-configured HTTP client to ensure the request doesn't hang
// indefinitely. It performs a standard GET request and reads the raw JSON
// response containing the public keys.
//
// The HTTP client is configured with JWKS_CLIENT_TIMEOUT (10 seconds by default)
// to prevent long-running requests. This timeout can be adjusted based on network
// conditions and reliability requirements.
//
// Parameters:
//   - jwksURL: The URL endpoint from which to fetch the JWKS
//
// Returns:
//   - string: The raw JSON string containing the JWKS data if successful
//   - error: An error if the HTTP request fails, the response cannot be read,
//     or any other error occurs during the fetching process
//
// This function is typically called by GetJWKS when the cached JWKS has expired
// or is not available, and should not be called directly in most cases.
func fetchJWKS(jwksURL string) (string, error) {
	return fetchJWKSWithTimeout(jwksURL, JWKS_CLIENT_TIMEOUT)
}

func fetchJWKSWithTimeout(jwksURL string, timeout time.Duration) (string, error) {
	client := &http.Client{Timeout: timeout}
	req, err := http.NewRequest("GET", jwksURL, nil)
	if err != nil {
		return "", err
	}

	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	jwksData, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	return string(jwksData), nil
}