package s3api

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"io"
	"testing"

	"github.com/seaweedfs/seaweedfs/weed/s3api/s3_constants"
)

// TestSSES3MultipartUploadStoresDerivedIV verifies the critical fix where
// handleSSES3MultipartEncryption must store the DERIVED IV (not base IV)
// in the returned key so it gets serialized into chunk metadata.
//
// This test prevents the bug where the derived IV was discarded, causing
// decryption to use the wrong IV and produce corrupted plaintext.
func TestSSES3MultipartUploadStoresDerivedIV(t *testing.T) {
	// Setup: Create a test key and base IV
	keyManager := GetSSES3KeyManager()
	sseS3Key, err := keyManager.GetOrCreateKey("")
	if err != nil {
		t.Fatalf("Failed to create SSE-S3 key: %v", err)
	}

	// Generate a random base IV
	baseIV := make([]byte, aes.BlockSize)
	if _, err := rand.Read(baseIV); err != nil {
		t.Fatalf("Failed to generate base IV: %v", err)
	}

	// Test data for multipart upload parts
	testCases := []struct {
		name       string
		partOffset int64
		data       []byte
	}{
		{"Part 1 at offset 0", 0, []byte("First part of multipart upload")},
		{"Part 2 at offset 1MB", 1024 * 1024, []byte("Second part of multipart upload")},
		{"Part 3 at offset 5MB", 5 * 1024 * 1024, []byte("Third part at 5MB offset")},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Calculate the expected derived IV (what encryption will use)
			expectedDerivedIV, ivSkip := calculateIVWithOffset(baseIV, tc.partOffset)

			// Call CreateSSES3EncryptedReaderWithBaseIV to encrypt the data
			dataReader := bytes.NewReader(tc.data)
			encryptedReader, returnedDerivedIV, encErr := CreateSSES3EncryptedReaderWithBaseIV(
				dataReader,
				sseS3Key,
				baseIV,
				tc.partOffset,
			)
			if encErr != nil {
				t.Fatalf("Failed to create encrypted reader: %v", encErr)
			}

			// Read the encrypted data
			encryptedData, err := io.ReadAll(encryptedReader)
			if err != nil {
				t.Fatalf("Failed to read encrypted data: %v", err)
			}

			// CRITICAL VERIFICATION: The returned IV should be the DERIVED IV
			if !bytes.Equal(returnedDerivedIV, expectedDerivedIV) {
				t.Errorf("CreateSSES3EncryptedReaderWithBaseIV returned wrong IV:\nExpected: %x\nGot: %x",
					expectedDerivedIV[:8], returnedDerivedIV[:8])
			}

			// CRITICAL TEST: Verify the key.IV field would be updated (simulating handleSSES3MultipartEncryption)
			// This is what the fix does: key.IV = derivedIV
			keyWithDerivedIV := &SSES3Key{
				Key:       sseS3Key.Key,
				KeyID:     sseS3Key.KeyID,
				Algorithm: sseS3Key.Algorithm,
				IV:        returnedDerivedIV, // This simulates: key.IV = derivedIV
			}

			// TEST 1: Verify decryption with DERIVED IV produces correct plaintext (correct behavior)
			decryptedWithDerivedIV := make([]byte, len(encryptedData))
			block, err := aes.NewCipher(keyWithDerivedIV.Key)
			if err != nil {
				t.Fatalf("Failed to create cipher: %v", err)
			}
			stream := cipher.NewCTR(block, keyWithDerivedIV.IV)

			// Handle ivSkip for non-block-aligned offsets
			if ivSkip > 0 {
				skipDummy := make([]byte, ivSkip)
				stream.XORKeyStream(skipDummy, skipDummy)
			}
			stream.XORKeyStream(decryptedWithDerivedIV, encryptedData)

			if !bytes.Equal(decryptedWithDerivedIV, tc.data) {
				t.Errorf("Decryption with derived IV failed:\nExpected: %q\nGot: %q",
					tc.data, decryptedWithDerivedIV)
			} else {
				t.Logf("✓ Derived IV decryption successful for offset %d", tc.partOffset)
			}

			// TEST 2: Verify decryption with BASE IV produces WRONG plaintext (bug behavior)
			// This is what would happen if the bug wasn't fixed
			if tc.partOffset > 0 { // Only test for non-zero offsets (where IVs differ)
				keyWithBaseIV := &SSES3Key{
					Key:       sseS3Key.Key,
					KeyID:     sseS3Key.KeyID,
					Algorithm: sseS3Key.Algorithm,
					IV:        baseIV, // BUG: Using base IV instead of derived IV
				}

				decryptedWithBaseIV := make([]byte, len(encryptedData))
				blockWrong, err := aes.NewCipher(keyWithBaseIV.Key)
				if err != nil {
					t.Fatalf("Failed to create cipher for wrong decryption: %v", err)
				}
				streamWrong := cipher.NewCTR(blockWrong, keyWithBaseIV.IV)
				streamWrong.XORKeyStream(decryptedWithBaseIV, encryptedData)

				if bytes.Equal(decryptedWithBaseIV, tc.data) {
					t.Errorf("CRITICAL BUG: Base IV produced correct plaintext at offset %d! Should produce corrupted data.", tc.partOffset)
				} else {
					t.Logf("✓ Verified: Base IV produces corrupted data at offset %d (bug would cause this)", tc.partOffset)
				}
			}
		})
	}
}

// TestHandleSSES3MultipartEncryptionFlow is an integration test that verifies
// the complete flow of handleSSES3MultipartEncryption, including that the
// returned key contains the derived IV (not base IV).
func TestHandleSSES3MultipartEncryptionFlow(t *testing.T) {
	// This test simulates what happens in a real multipart upload request

	// Generate test key manually (simulating a complete SSE-S3 key)
	keyBytes := make([]byte, 32) // 256-bit key
	if _, err := rand.Read(keyBytes); err != nil {
		t.Fatalf("Failed to generate key: %v", err)
	}

	originalKey := &SSES3Key{
		Key:       keyBytes,
		KeyID:     "test-key-id",
		Algorithm: SSES3Algorithm,
		IV:        nil, // Will be set later
	}

	baseIV := make([]byte, aes.BlockSize)
	if _, err := rand.Read(baseIV); err != nil {
		t.Fatalf("Failed to generate base IV: %v", err)
	}

	// For this test, we'll work directly with the key structure
	// since SerializeSSES3Metadata requires KMS setup

	// Test with a non-zero offset (where base IV != derived IV)
	partOffset := int64(2 * 1024 * 1024) // 2MB offset
	plaintext := []byte("Test data for part 2 of multipart upload")

	// Calculate what the derived IV should be
	expectedDerivedIV, ivSkip := calculateIVWithOffset(baseIV, partOffset)

	// Simulate the upload by calling CreateSSES3EncryptedReaderWithBaseIV directly
	// (This is what handleSSES3MultipartEncryption does internally)
	dataReader := bytes.NewReader(plaintext)

	// Encrypt with base IV and offset
	encryptedReader, derivedIV, encErr := CreateSSES3EncryptedReaderWithBaseIV(
		dataReader,
		originalKey,
		baseIV,
		partOffset,
	)
	if encErr != nil {
		t.Fatalf("Failed to create encrypted reader: %v", encErr)
	}

	// THE FIX: Update key.IV with derivedIV (this is what the bug fix does)
	originalKey.IV = derivedIV

	// Read encrypted data
	encryptedData, err := io.ReadAll(encryptedReader)
	if err != nil {
		t.Fatalf("Failed to read encrypted data: %v", err)
	}

	// VERIFICATION 1: Derived IV should match expected
	if !bytes.Equal(derivedIV, expectedDerivedIV) {
		t.Errorf("Derived IV mismatch:\nExpected: %x\nGot: %x",
			expectedDerivedIV[:8], derivedIV[:8])
	}

	// VERIFICATION 2: Key should now contain derived IV (the fix)
	if !bytes.Equal(originalKey.IV, derivedIV) {
		t.Errorf("Key.IV was not updated with derived IV!\nKey.IV: %x\nDerived IV: %x",
			originalKey.IV[:8], derivedIV[:8])
	} else {
		t.Logf("✓ Key.IV correctly updated with derived IV")
	}

	// VERIFICATION 3: The IV stored in the key can be used for decryption
	decryptedData := make([]byte, len(encryptedData))
	block, err := aes.NewCipher(originalKey.Key)
	if err != nil {
		t.Fatalf("Failed to create cipher: %v", err)
	}

	stream := cipher.NewCTR(block, originalKey.IV)

	// Handle ivSkip for non-block-aligned offsets
	if ivSkip > 0 {
		skipDummy := make([]byte, ivSkip)
		stream.XORKeyStream(skipDummy, skipDummy)
	}
	stream.XORKeyStream(decryptedData, encryptedData)

	if !bytes.Equal(decryptedData, plaintext) {
		t.Errorf("Final decryption failed:\nExpected: %q\nGot: %q", plaintext, decryptedData)
	} else {
		t.Logf("✓ Full encrypt-update_key-decrypt cycle successful")
	}
}

// TestSSES3HeaderEncoding tests that the header encoding/decoding works correctly
func TestSSES3HeaderEncoding(t *testing.T) {
	// Generate test base IV
	baseIV := make([]byte, aes.BlockSize)
	if _, err := rand.Read(baseIV); err != nil {
		t.Fatalf("Failed to generate base IV: %v", err)
	}

	// Encode as it would be in HTTP header
	baseIVHeader := base64.StdEncoding.EncodeToString(baseIV)

	// Decode (as handleSSES3MultipartEncryption does)
	decodedBaseIV, err := base64.StdEncoding.DecodeString(baseIVHeader)
	if err != nil {
		t.Fatalf("Failed to decode base IV: %v", err)
	}

	// Verify round-trip
	if !bytes.Equal(decodedBaseIV, baseIV) {
		t.Errorf("Base IV encoding round-trip failed:\nOriginal: %x\nDecoded: %x",
			baseIV, decodedBaseIV)
	}

	// Verify length
	if len(decodedBaseIV) != s3_constants.AESBlockSize {
		t.Errorf("Decoded base IV has wrong length: expected %d, got %d",
			s3_constants.AESBlockSize, len(decodedBaseIV))
	}
}