package cognitoauth

import (
	"crypto/rand"
	"crypto/rsa"
	"encoding/json"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
	"time"

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

// generateTestJWKS creates a test JSON Web Key Set (JWKS) for testing purposes.
// It performs the following operations:
// - Generates a 2048-bit RSA key pair for cryptographic operations
// - Creates a JWK (JSON Web Key) from the public key portion
// - Sets key properties including ID ("test-key-id"), algorithm (RS256), and usage ("sig")
// - Creates a JWK Set and adds the key to it
// - Returns both the JWK Set object and its JSON string representation
func generateTestJWKS(t *testing.T) (jwk.Set, string) {
	// Generate RSA key for testing
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		t.Fatalf("failed to generate RSA key: %v", err)
	}

	// Create JWK from RSA key
	key, err := jwk.FromRaw(privateKey.PublicKey)
	if err != nil {
		t.Fatalf("failed to create JWK: %v", err)
	}

	// Set key properties
	if err := key.Set(jwk.KeyIDKey, "test-key-id"); err != nil {
		t.Fatalf("failed to set key ID: %v", err)
	}
	if err := key.Set(jwk.AlgorithmKey, jwa.RS256); err != nil {
		t.Fatalf("failed to set algorithm: %v", err)
	}
	if err := key.Set(jwk.KeyUsageKey, "sig"); err != nil {
		t.Fatalf("failed to set key usage: %v", err)
	}

	// Create JWK Set with the key
	keySet := jwk.NewSet()
	if err := keySet.AddKey(key); err != nil {
		t.Fatalf("failed to add key to key set: %v", err)
	}

	// Serialize to JSON
	jwksJSON, err := json.Marshal(keySet)
	if err != nil {
		t.Fatalf("failed to marshal JWK Set: %v", err)
	}

	return keySet, string(jwksJSON)
}

// TestGetJWKS tests the GetJWKS function through multiple scenarios including
// initial fetching, caching, cache expiration, and error handling.
func TestGetJWKS(t *testing.T) {
	// Setup logger
	logHandler := slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{
		Level: slog.LevelDebug,
	})
	logger := slog.New(logHandler)

	// Generate test JWK set for mock responses
	testKeySet, jwksJSON := generateTestJWKS(t)

	// Mock server to return JWKS
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_, err := w.Write([]byte(jwksJSON))
		if err != nil {
			t.Errorf("Failed to write response: %v", err)
		}
	}))
	defer server.Close()

	// Reset JWKS cache before tests
	jwksCache = &JWKSCache{
		ExpiresAt: time.Now(), // Initial state is expired
	}

	// Test: First fetch - should make HTTP request
	// Tests initial JWKS retrieval and verifies:
	// - Successful HTTP request and response processing
	// - Correct number of keys in the retrieved set
	// - Proper caching of the retrieved key set
	// - Correct cache expiration time setting
	t.Run("first fetch", func(t *testing.T) {
		keySet, err := GetJWKS(server.URL, logger)
		if err != nil {
			t.Fatalf("GetJWKS() error = %v", err)
		}

		// Verify key set contains expected keys
		if keySet.Len() != testKeySet.Len() {
			t.Errorf("GetJWKS() returned %d keys, want %d", keySet.Len(), testKeySet.Len())
		}

		// Check that it's cached
		if jwksCache.KeySet == nil {
			t.Errorf("KeySet was not cached after fetch")
		}
		if time.Now().After(jwksCache.ExpiresAt) {
			t.Errorf("Cache expiration not set correctly")
		}
	})

	// Test: Subsequent fetch - should use cache
	// Verifies caching behavior:
	// - Tests that subsequent requests use cached data instead of making HTTP requests
	// - Tests the cache works even when the server is unavailable/failing
	// - Validates that the cached key count matches expected values
	t.Run("cached fetch", func(t *testing.T) {
		// Create a mock server that fails - to verify we're using the cache
		failServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusInternalServerError)
			_, err := w.Write([]byte("This server should not be called"))
			if err != nil {
				t.Errorf("Failed to write response: %v", err)
			}
		}))
		defer failServer.Close()

		// We should still get valid keys because we use the cache
		keySet, err := GetJWKS(failServer.URL, logger)
		if err != nil {
			t.Fatalf("GetJWKS() with cache error = %v", err)
		}

		// Verify key set contains expected keys
		if keySet.Len() != testKeySet.Len() {
			t.Errorf("GetJWKS() with cache returned %d keys, want %d", keySet.Len(), testKeySet.Len())
		}
	})

	// Test: Cache expiry - should fetch again
	// Tests cache expiration behavior:
	// - Verifies new keys are fetched after cache expiration
	// - Ensures the system properly handles refreshing expired cache data
	// - Validates that updated keys are received correctly
	t.Run("cache expiry", func(t *testing.T) {
		// Force cache expiration
		jwksCache.mutex.Lock()
		jwksCache.ExpiresAt = time.Now().Add(-1 * time.Hour) // Expired one hour ago
		jwksCache.mutex.Unlock()

		// New server with updated keys
		updatedKeySet, updatedJSON := generateTestJWKS(t)
		updateServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			_, err := w.Write([]byte(updatedJSON))
			if err != nil {
				t.Errorf("Failed to write response: %v", err)
			}
		}))
		defer updateServer.Close()

		// Now we should fetch new keys
		keySet, err := GetJWKS(updateServer.URL, logger)
		if err != nil {
			t.Fatalf("GetJWKS() after expiry error = %v", err)
		}

		// Verify key set contains expected keys (should match updated set)
		if keySet.Len() != updatedKeySet.Len() {
			t.Errorf("GetJWKS() after expiry returned %d keys, want %d", keySet.Len(), updatedKeySet.Len())
		}
	})

	// Test: Server error
	// Tests error handling when server fails:
	// - Validates behavior with 500 Internal Server Error responses
	// - Verifies appropriate error is returned to caller
	// - Tests system behavior with empty/reset cache
	t.Run("server error", func(t *testing.T) {
		// Reset cache
		jwksCache = &JWKSCache{
			ExpiresAt: time.Now(), // Initial state is expired
		}

		// Create a mock server that returns an error
		errorServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusInternalServerError)
			_, err := w.Write([]byte("Internal Server Error"))
			if err != nil {
				t.Errorf("Failed to write response: %v", err)
			}
		}))
		defer errorServer.Close()

		// Should return an error
		_, err := GetJWKS(errorServer.URL, logger)
		if err == nil {
			t.Errorf("GetJWKS() with server error should return error, got nil")
		}
	})

	// Test: Invalid JSON response
	// Tests handling of malformed responses:
	// - Verifies error handling for invalid JSON responses
	// - Ensures system stability when receiving corrupted data
	// - Validates error reporting for bad responses
	t.Run("invalid json", func(t *testing.T) {
		// Reset cache
		jwksCache = &JWKSCache{
			ExpiresAt: time.Now(), // Initial state is expired
		}

		// Create a mock server that returns invalid JSON
		invalidServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			_, err := w.Write([]byte("This is not valid JSON"))
			if err != nil {
				t.Errorf("Failed to write response: %v", err)
			}
		}))
		defer invalidServer.Close()

		// Should return an error
		_, err := GetJWKS(invalidServer.URL, logger)
		if err == nil {
			t.Errorf("GetJWKS() with invalid JSON should return error, got nil")
		}
	})
}

// TestFetchJWKS tests the fetchJWKS function directly with various scenarios
// including successful fetches, timeouts, invalid URLs, and error responses.
func TestFetchJWKS(t *testing.T) {
	// Generate test JWK set for mock responses
	_, jwksJSON := generateTestJWKS(t)

	// Test: Successful fetch
	// Verifies basic JWKS retrieval functionality:
	// - Tests successful HTTP request and response handling
	// - Validates that response body matches expected JSON
	// - Confirms proper return of fetched content
	t.Run("successful fetch", func(t *testing.T) {
		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Set("Content-Type", "application/json")
			_, err := w.Write([]byte(jwksJSON))
			if err != nil {
				t.Errorf("Failed to write response: %v", err)
			}
		}))
		defer server.Close()

		result, err := fetchJWKS(server.URL)
		if err != nil {
			t.Fatalf("fetchJWKS() error = %v", err)
		}

		if result != jwksJSON {
			t.Errorf("fetchJWKS() = %v, want %v", result, jwksJSON)
		}
	})

	// Test: Timeout
	// Validates timeout handling:
	// - Tests behavior when server response exceeds client timeout limit
	// - Verifies appropriate error handling for timed-out requests
	// - Ensures the function fails properly in timeout scenarios
	t.Run("timeout", func(t *testing.T) {
		t.Parallel()
		// Create a server that sleeps longer than the client timeout
		slowServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			time.Sleep(50 * time.Millisecond)
			_, err := w.Write([]byte(jwksJSON))
			if err != nil {
				t.Errorf("Failed to write response: %v", err)
			}
		}))
		defer slowServer.Close()

		_, err := fetchJWKSWithTimeout(slowServer.URL, 10*time.Millisecond)
		if err == nil {
			t.Errorf("fetchJWKS() with timeout should return error, got nil")
		}
	})

	// Test: Invalid URL
	// Tests error handling for bad URLs:
	// - Verifies behavior with non-existent or invalid URLs
	// - Ensures appropriate error responses for network failures
	// - Validates error handling for connection issues
	t.Run("invalid url", func(t *testing.T) {
		_, err := fetchJWKS("http://invalid-url-that-does-not-exist.example.com")
		if err == nil {
			t.Errorf("fetchJWKS() with invalid URL should return error, got nil")
		}
	})

	// Test: HTTP error status
	// Validates handling of HTTP error responses:
	// - Tests behavior with various HTTP error status codes (e.g., 404)
	// - Verifies response body handling regardless of status code
	// - Tests that function returns body content even for error status codes
	t.Run("http error status", func(t *testing.T) {
		errorServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(http.StatusNotFound)
			_, err := w.Write([]byte("Not Found"))
			if err != nil {
				t.Errorf("Failed to write response: %v", err)
			}
		}))
		defer errorServer.Close()

		result, err := fetchJWKS(errorServer.URL)
		if err != nil {
			t.Fatalf("fetchJWKS() unexpected error = %v", err)
		}

		// Function doesn't check status code, just reads body
		if result != "Not Found" {
			t.Errorf("fetchJWKS() = %v, want %v", result, "Not Found")
		}
	})
}