package awsauth

import (
	"context"
	"testing"
	"time"

	"github.com/aws/aws-sdk-go-v2/aws"
	"github.com/patrickmn/go-cache"
	"github.com/stretchr/testify/assert"
)

func TestGetSigV4Headers(t *testing.T) {
	ctx := context.Background()

	t.Run("success", func(t *testing.T) {
		// Clear any existing cache entries
		sigv4CredCache.Flush()

		authInfo := SigV4AuthInfo{
			RoleArn:    "arn:aws:iam::123456789012:role/valid-role",
			ExternalID: "valid-external-id",
			Region:     "us-west-2",
		}

		// Pre-populate cache with valid credentials to bypass assumeRole
		cacheKey := "arn:aws:iam::123456789012:role/valid-role:valid-external-id"
		creds := aws.Credentials{
			AccessKeyID:     "AKIAIOSFODNN7EXAMPLE",
			SecretAccessKey: "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY",
			SessionToken:    "session-token",
			Expires:         time.Now().Add(1 * time.Hour),
		}
		sigv4CredCache.Set(cacheKey, creds, cache.DefaultExpiration)

		headers, err := GetSigV4Headers(ctx, authInfo, "GET", "https://example.com/api", nil)
		assert.NoError(t, err)
		assert.NotNil(t, headers)
		assert.Contains(t, headers, "Authorization")
		assert.Contains(t, headers, "X-Amz-Date")
		assert.Contains(t, headers["Authorization"], "AWS4-HMAC-SHA256")
		assert.Contains(t, headers["Authorization"], "Credential=AKIAIOSFODNN7EXAMPLE")

		// Clean up
		sigv4CredCache.Flush()
	})

	t.Run("empty role ARN", func(t *testing.T) {
		authInfo := SigV4AuthInfo{
			RoleArn:    "",
			ExternalID: "",
			Region:     "us-west-2",
		}
		headers, err := GetSigV4Headers(ctx, authInfo, "GET", "https://example.com/api", nil)
		assert.Error(t, err)
		assert.Nil(t, headers)
		assert.Contains(t, err.Error(), "authInfo missing for SigV4 signing")
	})

	t.Run("empty region", func(t *testing.T) {
		authInfo := SigV4AuthInfo{
			RoleArn:    "arn:aws:iam::123456789012:role/test-role",
			ExternalID: "",
			Region:     "",
		}
		headers, err := GetSigV4Headers(ctx, authInfo, "GET", "https://example.com/api", nil)
		assert.Error(t, err)
		assert.Nil(t, headers)
		assert.Contains(t, err.Error(), "authInfo missing for SigV4 signing")
	})

	t.Run("empty role ARN and region", func(t *testing.T) {
		authInfo := SigV4AuthInfo{
			RoleArn:    "",
			ExternalID: "",
			Region:     "",
		}
		headers, err := GetSigV4Headers(ctx, authInfo, "GET", "https://example.com/api", nil)
		assert.Error(t, err)
		assert.Nil(t, headers)
	})
}

func TestSigV4CredCache(t *testing.T) {
	t.Run("cache operations", func(t *testing.T) {
		// Clear any existing cache entries
		sigv4CredCache.Flush()

		cacheKey := "test-role-arn:test-external-id"
		creds := aws.Credentials{
			AccessKeyID:     "test-access-key",
			SecretAccessKey: "test-secret-key",
			SessionToken:    "test-session-token",
			Expires:         time.Now().Add(1 * time.Hour),
		}

		// Set cache
		sigv4CredCache.Set(cacheKey, creds, cache.DefaultExpiration)

		// Get from cache
		cached, found := sigv4CredCache.Get(cacheKey)
		assert.True(t, found)
		assert.NotNil(t, cached)

		cachedCreds := cached.(aws.Credentials)
		assert.Equal(t, "test-access-key", cachedCreds.AccessKeyID)
		assert.Equal(t, "test-secret-key", cachedCreds.SecretAccessKey)
		assert.Equal(t, "test-session-token", cachedCreds.SessionToken)

		// Clean up
		sigv4CredCache.Flush()
	})

	t.Run("cache miss", func(t *testing.T) {
		sigv4CredCache.Flush()

		_, found := sigv4CredCache.Get("non-existent-key")
		assert.False(t, found)
	})
}

func TestGetOrAssumeRole_CacheHit(t *testing.T) {
	t.Run("returns cached credentials when valid", func(t *testing.T) {
		// Clear any existing cache entries
		sigv4CredCache.Flush()

		authInfo := SigV4AuthInfo{
			RoleArn:    "arn:aws:iam::123456789012:role/cached-role",
			ExternalID: "cached-external-id",
			Region:     "us-west-2",
		}

		cacheKey := "arn:aws:iam::123456789012:role/cached-role:cached-external-id"
		creds := aws.Credentials{
			AccessKeyID:     "cached-access-key",
			SecretAccessKey: "cached-secret-key",
			SessionToken:    "cached-session-token",
			Expires:         time.Now().Add(1 * time.Hour), // Valid for 1 hour
		}

		// Pre-populate cache
		sigv4CredCache.Set(cacheKey, creds, cache.DefaultExpiration)

		// Should return cached credentials
		result, err := getOrAssumeRole(context.Background(), authInfo)
		assert.NoError(t, err)
		assert.NotNil(t, result)
		assert.Equal(t, "cached-access-key", result.AccessKeyID)
		assert.Equal(t, "cached-secret-key", result.SecretAccessKey)
		assert.Equal(t, "cached-session-token", result.SessionToken)

		// Clean up
		sigv4CredCache.Flush()
	})

	t.Run("cache miss triggers assume role", func(t *testing.T) {
		sigv4CredCache.Flush()

		authInfo := SigV4AuthInfo{
			RoleArn:    "arn:aws:iam::123456789012:role/uncached-role",
			ExternalID: "",
			Region:     "us-west-2",
		}

		// This will fail because we don't have actual AWS credentials
		// but it verifies the flow doesn't return cached creds
		_, err := getOrAssumeRole(context.Background(), authInfo)
		assert.Error(t, err) // Expected to fail without real AWS creds
	})
}

func TestGetOrAssumeRole_ExpiredCache(t *testing.T) {
	t.Run("expired credentials trigger refresh", func(t *testing.T) {
		sigv4CredCache.Flush()

		authInfo := SigV4AuthInfo{
			RoleArn:    "arn:aws:iam::123456789012:role/expired-role",
			ExternalID: "expired-external-id",
			Region:     "us-west-2",
		}

		cacheKey := "arn:aws:iam::123456789012:role/expired-role:expired-external-id"
		creds := aws.Credentials{
			AccessKeyID:     "expired-access-key",
			SecretAccessKey: "expired-secret-key",
			SessionToken:    "expired-session-token",
			Expires:         time.Now().Add(-1 * time.Hour), // Expired 1 hour ago
			CanExpire:       true,
		}

		// Pre-populate cache with expired credentials
		sigv4CredCache.Set(cacheKey, creds, cache.DefaultExpiration)

		// Should attempt to refresh since credentials are expired
		// This will fail because we don't have actual AWS credentials
		_, err := getOrAssumeRole(context.Background(), authInfo)
		assert.Error(t, err) // Expected to fail without real AWS creds

		// Clean up
		sigv4CredCache.Flush()
	})
}

func TestAssumeRole_InvalidCredentials(t *testing.T) {
	t.Run("fails without valid AWS credentials", func(t *testing.T) {
		authInfo := SigV4AuthInfo{
			RoleArn:    "arn:aws:iam::123456789012:role/test-role",
			ExternalID: "test-external-id",
			Region:     "us-west-2",
		}

		// This should fail because we don't have valid AWS credentials
		_, err := assumeRole(context.Background(), authInfo)
		assert.Error(t, err)
	})

	t.Run("uses default region when not specified", func(t *testing.T) {
		t.Setenv("AWS_REGION", "eu-west-1")

		authInfo := SigV4AuthInfo{
			RoleArn:    "arn:aws:iam::123456789012:role/test-role",
			ExternalID: "",
			Region:     "", // Empty region should use env var
		}

		// This will fail but we're testing that region fallback works
		_, err := assumeRole(context.Background(), authInfo)
		assert.Error(t, err) // Expected to fail without real AWS creds
	})
}