package jwt_utils

import (
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

var testSecret = []byte("test-secret-key")

func TestGenerateJWTToken(t *testing.T) {
	uuid := "test-uuid"
	duration := time.Hour

	token, expiry, err := GenerateJWTToken(testSecret, uuid, duration)

	assert.NoError(t, err)
	assert.NotEmpty(t, token)
	assert.True(t, expiry.After(time.Now()))
	assert.True(t, expiry.Before(time.Now().Add(duration).Add(time.Second)))

	// Verify token can be processed
	processedUUID, err := ProcessJWTToken(testSecret, token)
	assert.NoError(t, err)
	assert.Equal(t, uuid, processedUUID)
}

func TestGenerateJWTTokenPair(t *testing.T) {
	uuid := "test-uuid"
	accessDuration := time.Hour
	refreshDuration := time.Hour * 24

	accessToken, accessExpiry, refreshToken, refreshExpiry, err := GenerateJWTTokenPair(
		testSecret,
		uuid,
		accessDuration,
		refreshDuration,
	)

	assert.NoError(t, err)
	assert.NotEmpty(t, accessToken)
	assert.NotEmpty(t, refreshToken)
	assert.True(t, accessExpiry.After(time.Now()))
	assert.True(t, refreshExpiry.After(time.Now()))
	assert.True(t, accessExpiry.Before(time.Now().Add(accessDuration).Add(time.Second)))
	assert.True(t, refreshExpiry.Before(time.Now().Add(refreshDuration).Add(time.Second)))

	// Verify both tokens can be processed
	processedAccessUUID, err := ProcessJWTToken(testSecret, accessToken)
	assert.NoError(t, err)
	assert.Equal(t, uuid, processedAccessUUID)

	processedRefreshUUID, err := ProcessJWTToken(testSecret, refreshToken)
	assert.NoError(t, err)
	assert.Equal(t, uuid, processedRefreshUUID)
}

func TestProcessJWTToken_Invalid(t *testing.T) {
	tests := []struct {
		name    string
		token   string
		wantErr bool
	}{
		{
			name:    "empty token",
			token:   "",
			wantErr: true,
		},
		{
			name:    "malformed token",
			token:   "not.a.token",
			wantErr: true,
		},
		{
			name:    "wrong signature",
			token:   "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzZXNzaW9uX3V1aWQiOiJ0ZXN0LXV1aWQiLCJleHAiOjE3MDQwNjc1NTF9.wrong",
			wantErr: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			uuid, err := ProcessJWTToken(testSecret, tt.token)
			if tt.wantErr {
				assert.Error(t, err)
				assert.Empty(t, uuid)
			} else {
				assert.NoError(t, err)
				assert.NotEmpty(t, uuid)
			}
		})
	}
}

func TestProcessJWTToken_Expired(t *testing.T) {
	uuid := "test-uuid"
	duration := -time.Hour // negative duration for expired token

	token, _, err := GenerateJWTToken(testSecret, uuid, duration)
	assert.NoError(t, err)

	processedUUID, err := ProcessJWTToken(testSecret, token)
	assert.Error(t, err)
	assert.Empty(t, processedUUID)
}

// TestProcessJWTToken_AlgorithmConfusion tests protection against JWT algorithm confusion attacks
// CVE-2015-9235: None algorithm bypass
// CVE-2016-5431: HS256/RS256 algorithm confusion
// CWE-347: Improper Verification of Cryptographic Signature
func TestProcessJWTToken_AlgorithmConfusion(t *testing.T) {
	tests := []struct {
		name        string
		token       string
		description string
		wantErr     bool
	}{
		{
			name: "none algorithm bypass attempt",
			// Token with "alg": "none" - should be rejected
			token:       "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0.eyJzZXNzaW9uX3V1aWQiOiJhdHRhY2tlci11dWlkIiwiZXhwIjo5OTk5OTk5OTk5fQ.",
			description: "Attacker tries to bypass signature verification using 'none' algorithm",
			wantErr:     true,
		},
		{
			name: "RS256 algorithm confusion attempt",
			// Token with "alg": "RS256" - should be rejected (we only accept HS256)
			token:       "eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJzZXNzaW9uX3V1aWQiOiJhdHRhY2tlci11dWlkIiwiZXhwIjo5OTk5OTk5OTk5fQ.invalid",
			description: "Attacker tries to use RS256 to confuse HMAC validation",
			wantErr:     true,
		},
		{
			name: "HS384 algorithm attempt",
			// Token with "alg": "HS384" - should be rejected (we only accept HS256)
			token:       "eyJhbGciOiJIUzM4NCIsInR5cCI6IkpXVCJ9.eyJzZXNzaW9uX3V1aWQiOiJhdHRhY2tlci11dWlkIiwiZXhwIjo5OTk5OTk5OTk5fQ.invalid",
			description: "Attacker tries to use different HMAC algorithm",
			wantErr:     true,
		},
		{
			name: "HS512 algorithm attempt",
			// Token with "alg": "HS512" - should be rejected (we only accept HS256)
			token:       "eyJhbGciOiJIUzUxMiIsInR5cCI6IkpXVCJ9.eyJzZXNzaW9uX3V1aWQiOiJhdHRhY2tlci11dWlkIiwiZXhwIjo5OTk5OTk5OTk5fQ.invalid",
			description: "Attacker tries to use different HMAC algorithm",
			wantErr:     true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Logf("Testing: %s", tt.description)
			uuid, err := ProcessJWTToken(testSecret, tt.token)

			if tt.wantErr {
				assert.Error(t, err, "Expected error for security vulnerability: %s", tt.description)
				assert.Empty(t, uuid, "UUID should be empty when algorithm validation fails")
			} else {
				assert.NoError(t, err)
				assert.NotEmpty(t, uuid)
			}
		})
	}
}

// TestProcessJWTToken_ValidHS256Only tests that only valid HS256 tokens are accepted
func TestProcessJWTToken_ValidHS256Only(t *testing.T) {
	uuid := "valid-test-uuid"
	duration := time.Hour

	// Generate a valid HS256 token
	token, _, err := GenerateJWTToken(testSecret, uuid, duration)
	assert.NoError(t, err, "Should generate valid token")

	// Verify it's accepted
	processedUUID, err := ProcessJWTToken(testSecret, token)
	assert.NoError(t, err, "Valid HS256 token should be accepted")
	assert.Equal(t, uuid, processedUUID, "UUID should match")
}

// TestProcessJWTToken_MissingSessionUUID tests protection against missing session_uuid claim
func TestProcessJWTToken_MissingSessionUUID(t *testing.T) {
	// This test verifies the safe type assertion fix for CWE-391
	// A token without session_uuid claim should return an error, not panic

	// Note: We can't easily create such a token with our GenerateJWTToken function
	// as it always includes session_uuid. In a real attack scenario, an attacker
	// would craft such a token manually. This test documents the expected behavior.

	// For now, we verify that a malformed token is properly rejected
	malformedToken := "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJleHAiOjk5OTk5OTk5OTl9.invalid"
	uuid, err := ProcessJWTToken(testSecret, malformedToken)
	assert.Error(t, err, "Token without session_uuid should be rejected")
	assert.Empty(t, uuid, "UUID should be empty for invalid token")
}
