package s3api

import (
	"context"
	"fmt"
	"net/http"
	"net/url"
	"testing"
	"time"

	"github.com/seaweedfs/seaweedfs/weed/iam/integration"
	"github.com/seaweedfs/seaweedfs/weed/iam/policy"
	"github.com/seaweedfs/seaweedfs/weed/iam/sts"
	"github.com/seaweedfs/seaweedfs/weed/s3api/s3_constants"
	"github.com/seaweedfs/seaweedfs/weed/s3api/s3err"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestAssumeRole_CallerIdentityFallback tests the fallback logic when RoleArn is missing
func TestAssumeRole_CallerIdentityFallback(t *testing.T) {
	// Setup STS service
	stsService, _ := setupTestSTSService(t)

	// Create IAM integration mock
	iamMock := &MockIAMIntegration{
		authorizeFunc: func(ctx context.Context, identity *IAMIdentity, action Action, bucket, object string, r *http.Request) s3err.ErrorCode {
			// Allow global sts:AssumeRole
			if action == "sts:AssumeRole" {
				return s3err.ErrNone
			}
			return s3err.ErrAccessDenied
		},
		validateTrustPolicyFunc: func(ctx context.Context, roleArn, principalArn string) error {
			// Allow all trust policies for this test
			return nil
		},
	}

	// Create IAM service with the mock integration
	iam := &IdentityAccessManagement{
		iamIntegration: iamMock,
	}

	// Create STS handlers
	stsHandlers := NewSTSHandlers(stsService, iam)

	// Test case 1: Caller is an IAM User, RoleArn is missing
	t.Run("Caller is IAM User, No RoleArn", func(t *testing.T) {
		// Mock request
		req, err := http.NewRequest("POST", "/", nil)
		require.NoError(t, err)
		req.Form = url.Values{}
		req.Form.Set("Action", "AssumeRole")
		req.Form.Set("RoleSessionName", "test-session")
		req.Form.Set("Version", "2011-06-15")

		// Mock the authenticated identity (IAM User)
		callerIdentity := &Identity{
			Name:         "alice",
			Account:      &AccountAdmin,
			PrincipalArn: fmt.Sprintf("arn:aws:iam::%s:user/alice", defaultAccountID),
			Actions:      []Action{s3_constants.ACTION_ADMIN},
		}

		// 1. Test prepareSTSCredentials with NO RoleArn (simulating the fallback logic having passed PrincipalArn)
		// expected RoleArn passed to prepareSTSCredentials would be the caller's PrincipalArn
		fallbackRoleArn := callerIdentity.PrincipalArn

		// Prepare custom claims for the session (mimicking handleAssumeRole logic)
		var modifyClaims func(claims *sts.STSSessionClaims)
		if callerIdentity.isAdmin() {
			modifyClaims = func(claims *sts.STSSessionClaims) {
				if claims.RequestContext == nil {
					claims.RequestContext = make(map[string]interface{})
				}
				claims.RequestContext["is_admin"] = true
			}
		}

		stsCreds, assumedUser, err := stsHandlers.prepareSTSCredentials(context.Background(), fallbackRoleArn, "", "test-session", nil, "", modifyClaims)
		require.NoError(t, err)

		// Assertions
		// The role name should be extracted from the user ARN ("alice")
		assert.Contains(t, assumedUser.Arn, fmt.Sprintf("assumed-role/alice/test-session"))
		assert.Contains(t, assumedUser.AssumedRoleId, "alice:test-session")

		// Verify token claims using ValidateSessionToken
		sessionInfo, err := stsService.ValidateSessionToken(context.Background(), stsCreds.SessionToken)
		require.NoError(t, err)

		// The RoleArn in session info should match the fallback ARN (user ARN)
		assert.Equal(t, fallbackRoleArn, sessionInfo.RoleArn)

		// Verify is_admin claim is present
		isAdmin, ok := sessionInfo.RequestContext["is_admin"].(bool)
		assert.True(t, ok, "is_admin claim should be present")
		assert.True(t, isAdmin, "is_admin claim should be true")
	})

	// Test case 2: Caller is an STS Assumed Role, No RoleArn
	t.Run("Caller is STS Assumed Role, No RoleArn", func(t *testing.T) {
		// Mock identity
		callerIdentity := &Identity{
			Name:         "arn:aws:sts::111122223333:assumed-role/admin/session1",
			Account:      &AccountAdmin,
			PrincipalArn: "arn:aws:sts::111122223333:assumed-role/admin/session1",
		}

		fallbackRoleArn := callerIdentity.PrincipalArn

		stsCreds, assumedUser, err := stsHandlers.prepareSTSCredentials(context.Background(), fallbackRoleArn, "", "nested-session", nil, "", nil)
		require.NoError(t, err)

		// The role name should be extracted from the assumed role ARN ("admin")
		assert.Contains(t, assumedUser.Arn, "assumed-role/admin/nested-session")
		assert.Contains(t, assumedUser.AssumedRoleId, "admin:nested-session")

		// Check claims
		sessionInfo, err := stsService.ValidateSessionToken(context.Background(), stsCreds.SessionToken)
		require.NoError(t, err)
		assert.Equal(t, fallbackRoleArn, sessionInfo.RoleArn)
	})

	// Test case 3: Explicit RoleArn provided (Standard AssumeRole)
	t.Run("Explicit RoleArn Provided", func(t *testing.T) {
		explicitRoleArn := "arn:aws:iam::111122223333:role/TargetRole"

		stsCreds, assumedUser, err := stsHandlers.prepareSTSCredentials(context.Background(), explicitRoleArn, "", "explicit-session", nil, "", nil)
		require.NoError(t, err)

		// Role name should be "TargetRole"
		assert.Contains(t, assumedUser.Arn, "assumed-role/TargetRole/explicit-session")

		// Check claims
		sessionInfo, err := stsService.ValidateSessionToken(context.Background(), stsCreds.SessionToken)
		require.NoError(t, err)
		assert.Equal(t, explicitRoleArn, sessionInfo.RoleArn)
	})

	// Test case 4: Malformed ARN (Edge case)
	t.Run("Malformed ARN", func(t *testing.T) {
		malformedArn := "invalid-arn"

		stsCreds, assumedUser, err := stsHandlers.prepareSTSCredentials(context.Background(), malformedArn, "", "bad-session", nil, "", nil)
		require.NoError(t, err)

		// Fallback behavior: use full string as role name if extraction fails
		assert.Contains(t, assumedUser.Arn, "assumed-role/invalid-arn/bad-session")

		sessionInfo, err := stsService.ValidateSessionToken(context.Background(), stsCreds.SessionToken)
		require.NoError(t, err)
		assert.Equal(t, malformedArn, sessionInfo.RoleArn)
	})
}

func TestAssumeRole_EmbedsRolePolicies(t *testing.T) {
	t.Run("RoleWithAttachedPolicies", func(t *testing.T) {
		ctx := context.Background()
		manager := newTestSTSIntegrationManager(t)

		writePolicy := &policy.PolicyDocument{
			Version: "2012-10-17",
			Statement: []policy.Statement{
				{
					Effect: "Allow",
					Action: []string{"s3:*"},
					Resource: []string{
						"arn:aws:s3:::*",
						"arn:aws:s3:::*/*",
					},
				},
			},
		}
		require.NoError(t, manager.CreatePolicy(ctx, "", "S3WritePolicy", writePolicy))

		roleName := "LakekeeperVendedRole"
		require.NoError(t, manager.CreateRole(ctx, "", roleName, &integration.RoleDefinition{
			RoleName:         roleName,
			AttachedPolicies: []string{"S3WritePolicy"},
		}))

		iam := &IdentityAccessManagement{
			iamIntegration: NewS3IAMIntegration(manager, ""),
		}
		stsHandlers := NewSTSHandlers(manager.GetSTSService(), iam)

		roleArn := fmt.Sprintf("arn:aws:iam::%s:role/%s", defaultAccountID, roleName)
		stsCreds, _, err := stsHandlers.prepareSTSCredentials(ctx, roleArn, "", "test-session", nil, "", nil)
		require.NoError(t, err)

		sessionInfo, err := manager.GetSTSService().ValidateSessionToken(ctx, stsCreds.SessionToken)
		require.NoError(t, err)
		require.NotNil(t, sessionInfo)
		assert.Equal(t, []string{"S3WritePolicy"}, sessionInfo.Policies)
	})

	t.Run("RoleWithoutAttachedPolicies", func(t *testing.T) {
		ctx := context.Background()
		manager := newTestSTSIntegrationManager(t)

		roleName := "LakekeeperEmptyRole"
		require.NoError(t, manager.CreateRole(ctx, "", roleName, &integration.RoleDefinition{
			RoleName: roleName,
		}))

		iam := &IdentityAccessManagement{
			iamIntegration: NewS3IAMIntegration(manager, ""),
		}
		stsHandlers := NewSTSHandlers(manager.GetSTSService(), iam)

		roleArn := fmt.Sprintf("arn:aws:iam::%s:role/%s", defaultAccountID, roleName)
		stsCreds, _, err := stsHandlers.prepareSTSCredentials(ctx, roleArn, "", "test-session", nil, "", nil)
		require.NoError(t, err)

		sessionInfo, err := manager.GetSTSService().ValidateSessionToken(ctx, stsCreds.SessionToken)
		require.NoError(t, err)
		assert.Empty(t, sessionInfo.Policies)
	})
}

func newTestSTSIntegrationManager(t *testing.T) *integration.IAMManager {
	t.Helper()
	manager := integration.NewIAMManager()
	config := &integration.IAMConfig{
		STS: &sts.STSConfig{
			TokenDuration:    sts.FlexibleDuration{Duration: time.Hour},
			MaxSessionLength: sts.FlexibleDuration{Duration: 12 * time.Hour},
			Issuer:           "test-issuer",
			SigningKey:       []byte("test-signing-key-at-least-32-bytes-long-for-security"),
		},
		Policy: &policy.PolicyEngineConfig{
			DefaultEffect: "Deny",
			StoreType:     "memory",
		},
		Roles: &integration.RoleStoreConfig{
			StoreType: "memory",
		},
	}
	require.NoError(t, manager.Initialize(config, func() string { return "" }))
	return manager
}

// AssumeRole derives its session length from the same default-then-cap rule as
// the service layer (#11473): an omitted DurationSeconds yields TokenDuration,
// and either path is capped at MaxSessionLength.
func TestPrepareSTSCredentialsHonorsConfiguredDurations(t *testing.T) {
	stsService := sts.NewSTSService()
	require.NoError(t, stsService.Initialize(&sts.STSConfig{
		TokenDuration:    sts.FlexibleDuration{Duration: 15 * time.Minute},
		MaxSessionLength: sts.FlexibleDuration{Duration: 20 * time.Minute},
		Issuer:           "test-issuer",
		SigningKey:       []byte("test-signing-key-at-least-32-bytes-long-for-security"),
	}))
	stsHandlers := NewSTSHandlers(stsService, nil)
	roleArn := fmt.Sprintf("arn:aws:iam::%s:role/test-role", defaultAccountID)

	expiresIn := func(durationSeconds *int64) time.Duration {
		stsCreds, _, err := stsHandlers.prepareSTSCredentials(context.Background(), roleArn, "", "test-session", durationSeconds, "", nil)
		require.NoError(t, err)
		exp, err := time.Parse(time.RFC3339, stsCreds.Expiration)
		require.NoError(t, err)
		return time.Until(exp)
	}

	oneHour := int64(3600)
	assert.InDelta(t, (15 * time.Minute).Seconds(), expiresIn(nil).Seconds(), 60)
	assert.InDelta(t, (20 * time.Minute).Seconds(), expiresIn(&oneHour).Seconds(), 60)
}

// A named role's MaxSessionDuration bounds the session however DurationSeconds
// was resolved, matching the SDK paths' capDurationByRole.
func TestPrepareSTSCredentialsCapsAtRoleMaxDuration(t *testing.T) {
	ctx := context.Background()
	manager := newTestSTSIntegrationManager(t)
	require.NoError(t, manager.CreateRole(ctx, "", "ShortLivedRole", &integration.RoleDefinition{
		RoleName:           "ShortLivedRole",
		MaxSessionDuration: 3600,
	}))

	stsService := sts.NewSTSService()
	require.NoError(t, stsService.Initialize(&sts.STSConfig{
		TokenDuration:    sts.FlexibleDuration{Duration: 2 * time.Hour},
		MaxSessionLength: sts.FlexibleDuration{Duration: 12 * time.Hour},
		Issuer:           "test-issuer",
		SigningKey:       []byte("test-signing-key-at-least-32-bytes-long-for-security"),
	}))
	iam := &IdentityAccessManagement{iamIntegration: NewS3IAMIntegration(manager, "")}
	stsHandlers := NewSTSHandlers(stsService, iam)
	roleArn := fmt.Sprintf("arn:aws:iam::%s:role/ShortLivedRole", defaultAccountID)

	expiresIn := func(durationSeconds *int64) time.Duration {
		stsCreds, _, err := stsHandlers.prepareSTSCredentials(ctx, roleArn, "", "test-session", durationSeconds, "", nil)
		require.NoError(t, err)
		exp, err := time.Parse(time.RFC3339, stsCreds.Expiration)
		require.NoError(t, err)
		return time.Until(exp)
	}

	twoHours := int64(7200)
	assert.InDelta(t, float64(3600), expiresIn(nil).Seconds(), 60, "omitted duration resolves to the 2h default but the role caps it at 1h")
	assert.InDelta(t, float64(3600), expiresIn(&twoHours).Seconds(), 60, "explicit duration above the role max is capped")
}

// A session is issued from the role definition whose trust admits the
// caller. A role replaced after the caller's trust check by one that does
// not trust the caller yields no session, not one bound to the replacement.
func TestPrepareSTSCredentialsChecksTrustOnTheRoleItBinds(t *testing.T) {
	ctx := context.Background()
	manager := newTestSTSIntegrationManager(t)
	trusting := func(principal string) *policy.PolicyDocument {
		return &policy.PolicyDocument{Version: "2012-10-17", Statement: []policy.Statement{{
			Effect: "Allow", Action: []string{"sts:AssumeRole"},
			Principal: map[string]interface{}{"AWS": principal},
		}}}
	}
	caller := fmt.Sprintf("arn:aws:iam::%s:user/alice", defaultAccountID)
	roleName := "ReplacedRole"
	roleArn := fmt.Sprintf("arn:aws:iam::%s:role/%s", defaultAccountID, roleName)
	stsHandlers := NewSTSHandlers(manager.GetSTSService(), &IdentityAccessManagement{iamIntegration: NewS3IAMIntegration(manager, "")})

	// The caller's trust check passed against a role since replaced by one
	// trusting someone else.
	require.NoError(t, manager.CreateRole(ctx, "", roleName, &integration.RoleDefinition{
		RoleName: roleName, TrustPolicy: trusting(fmt.Sprintf("arn:aws:iam::%s:user/bob", defaultAccountID)),
	}))
	_, _, err := stsHandlers.prepareSTSCredentials(ctx, roleArn, caller, "s", nil, "", nil)
	require.ErrorIs(t, err, integration.ErrTrustPolicyDenied, "a session was issued for a role that does not trust the caller")

	// A replacement that does trust the caller binds its own ID.
	require.NoError(t, manager.CreateRole(ctx, "", roleName, &integration.RoleDefinition{RoleName: roleName, TrustPolicy: trusting(caller)}))
	role, err := manager.GetRole(ctx, roleName)
	require.NoError(t, err)
	creds, _, err := stsHandlers.prepareSTSCredentials(ctx, roleArn, caller, "s", nil, "", nil)
	require.NoError(t, err)
	session, err := manager.GetSTSService().ValidateSessionToken(ctx, creds.SessionToken)
	require.NoError(t, err)
	assert.Equal(t, role.RoleId, session.RoleId)
}