// Copyright (c) 2026 Alibaba Group and its affiliates

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package utils

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"testing"

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

func TestConvertPubKeyToPEM(t *testing.T) {
	// Generate a test RSA key pair
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	assert.NoError(t, err)

	// Test converting public key to PEM
	pubKeyPEM, err := ConvertRSAPubKeyToPEM(&privateKey.PublicKey)
	assert.NoError(t, err)
	assert.NotEmpty(t, pubKeyPEM)
	assert.Contains(t, pubKeyPEM, "-----BEGIN PUBLIC KEY-----")
	assert.Contains(t, pubKeyPEM, "-----END PUBLIC KEY-----")
}

func TestConvertPubKeyToPEMBytes(t *testing.T) {
	// Generate a test RSA key pair
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	assert.NoError(t, err)

	// Test converting public key to PEM bytes
	pubKeyPEMBytes, err := ConvertRSAPubKeyToPEMBytes(&privateKey.PublicKey)
	assert.NoError(t, err)
	assert.NotEmpty(t, pubKeyPEMBytes)
	assert.Contains(t, string(pubKeyPEMBytes), "-----BEGIN PUBLIC KEY-----")
	assert.Contains(t, string(pubKeyPEMBytes), "-----END PUBLIC KEY-----")
}

func TestConvertPubKeyToPEMInvalidKey(t *testing.T) {
	// Test with nil public key (this would cause a panic in the actual implementation)
	// We can't test this directly as it would panic, but we can note that this is a potential issue
}

func TestEncryptWithAESGCM(t *testing.T) {
	// Test with valid 32-byte session key
	sessionKey := make([]byte, 32)
	_, err := rand.Read(sessionKey)
	assert.NoError(t, err)

	plaintext := []byte("Hello, World! This is a test message.")
	ciphertext, err := EncryptWithAESGCM(plaintext, sessionKey)
	assert.NoError(t, err)
	assert.NotEmpty(t, ciphertext)
	assert.NotEqual(t, plaintext, ciphertext)

	// Test with invalid session key length
	invalidKey := make([]byte, 16) // 16 bytes instead of 32
	_, err = EncryptWithAESGCM(plaintext, invalidKey)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "invalid session key length")
}

func TestDecryptWithAESGCM(t *testing.T) {
	// Test with valid 32-byte session key
	sessionKey := make([]byte, 32)
	_, err := rand.Read(sessionKey)
	assert.NoError(t, err)

	plaintext := []byte("Hello, World! This is a test message.")
	ciphertext, err := EncryptWithAESGCM(plaintext, sessionKey)
	assert.NoError(t, err)

	// Decrypt the ciphertext
	decrypted, err := DecryptWithAESGCM(ciphertext, sessionKey)
	assert.NoError(t, err)
	assert.Equal(t, plaintext, decrypted)

	// Test with invalid session key length
	invalidKey := make([]byte, 16) // 16 bytes instead of 32
	_, err = DecryptWithAESGCM(ciphertext, invalidKey)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "invalid session key length")

	// Test with corrupted ciphertext
	corruptedCiphertext := make([]byte, len(ciphertext))
	copy(corruptedCiphertext, ciphertext)
	corruptedCiphertext[len(corruptedCiphertext)-1] ^= 0xFF // Flip last bit
	_, err = DecryptWithAESGCM(corruptedCiphertext, sessionKey)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "failed to decrypt data")

	// Test with too short ciphertext
	shortCiphertext := make([]byte, 10)
	_, err = DecryptWithAESGCM(shortCiphertext, sessionKey)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "ciphertext too short")
}

func TestParsePubKeyFromPEM(t *testing.T) {
	// Generate a test RSA key pair
	privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
	assert.NoError(t, err)

	// Convert the public key to PEM format
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	assert.NoError(t, err)
	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubKeyBytes,
	})

	// Test parsing valid PEM
	parsedPubKey, err := ParseRSAPubKeyFromPEM(string(pubKeyPEM))
	assert.NoError(t, err)
	assert.NotNil(t, parsedPubKey)
	assert.Equal(t, &privateKey.PublicKey, parsedPubKey)
}

func TestParsePubKeyFromPEMInvalidPEM(t *testing.T) {
	// Test parsing invalid PEM
	parsedPubKey, err := ParseRSAPubKeyFromPEM("invalid-pem-data")
	assert.Error(t, err)
	assert.Nil(t, parsedPubKey)
}

func TestParsePubKeyFromPEMInvalidKey(t *testing.T) {
	// Create invalid key data
	invalidKeyData := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: []byte("invalid key data"),
	})

	// Test parsing PEM with invalid key data
	parsedPubKey, err := ParseRSAPubKeyFromPEM(string(invalidKeyData))
	assert.Error(t, err)
	assert.Nil(t, parsedPubKey)
}

// === ECDSA/ECDH Tests ===

func TestConvertECDSAPubKeyToPEM(t *testing.T) {
	// Generate a test ECDSA key pair
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	assert.NoError(t, err)

	// Test converting public key to PEM string
	pubKeyPEM, err := ConvertECDSAPubKeyToPEM(&privateKey.PublicKey)
	assert.NoError(t, err)
	assert.NotEmpty(t, pubKeyPEM)
	assert.Contains(t, pubKeyPEM, "-----BEGIN PUBLIC KEY-----")
	assert.Contains(t, pubKeyPEM, "-----END PUBLIC KEY-----")
}

func TestConvertECDSAPubKeyToPEMBytes(t *testing.T) {
	// Generate a test ECDSA key pair
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	assert.NoError(t, err)

	// Test converting public key to PEM bytes
	pubKeyPEMBytes, err := ConvertECDSAPubKeyToPEMBytes(&privateKey.PublicKey)
	assert.NoError(t, err)
	assert.NotEmpty(t, pubKeyPEMBytes)
	assert.Contains(t, string(pubKeyPEMBytes), "-----BEGIN PUBLIC KEY-----")
	assert.Contains(t, string(pubKeyPEMBytes), "-----END PUBLIC KEY-----")
}

func TestParseECDSAPubKeyFromPEM(t *testing.T) {
	// Generate a test ECDSA key pair
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	assert.NoError(t, err)

	// Convert to PEM format
	pubKeyPEM, err := ConvertECDSAPubKeyToPEM(&privateKey.PublicKey)
	assert.NoError(t, err)

	// Test parsing valid PEM
	parsedPubKey, err := ParseECDSAPubKeyFromPEM(pubKeyPEM)
	assert.NoError(t, err)
	assert.NotNil(t, parsedPubKey)

	// Verify the parsed key matches the original
	assert.Equal(t, privateKey.PublicKey.X.Cmp(parsedPubKey.X), 0)
	assert.Equal(t, privateKey.PublicKey.Y.Cmp(parsedPubKey.Y), 0)
}

func TestParseECDSAPubKeyFromPEMInvalidPEM(t *testing.T) {
	// Test parsing invalid PEM
	parsedPubKey, err := ParseECDSAPubKeyFromPEM("invalid-pem-data")
	assert.Error(t, err)
	assert.Nil(t, parsedPubKey)
}

func TestMarshalECDHPublicKey(t *testing.T) {
	// Generate a test ECDSA key pair
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	assert.NoError(t, err)

	// Test marshaling public key to raw format
	pubKeyBytes := MarshalECDHPublicKey(&privateKey.PublicKey)
	assert.NotNil(t, pubKeyBytes)
	assert.Equal(t, 65, len(pubKeyBytes))       // 1 + 32 + 32 bytes for P-256
	assert.Equal(t, byte(0x04), pubKeyBytes[0]) // Uncompressed format prefix
}

func TestComputeECDHSharedSecret(t *testing.T) {
	// Generate two ECDSA key pairs
	alicePrivKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	assert.NoError(t, err)
	bobPrivKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	assert.NoError(t, err)

	// Get public key bytes
	alicePubKeyBytes := MarshalECDHPublicKey(&alicePrivKey.PublicKey)
	bobPubKeyBytes := MarshalECDHPublicKey(&bobPrivKey.PublicKey)

	// Compute shared secrets
	sharedSecret1, err := ComputeECDHSharedSecret(alicePrivKey, bobPubKeyBytes)
	assert.NoError(t, err)
	assert.NotNil(t, sharedSecret1)
	assert.Equal(t, 32, len(sharedSecret1)) // P-256 shared secret is 32 bytes

	sharedSecret2, err := ComputeECDHSharedSecret(bobPrivKey, alicePubKeyBytes)
	assert.NoError(t, err)
	assert.NotNil(t, sharedSecret2)
	assert.Equal(t, 32, len(sharedSecret2))

	// Verify ECDH symmetry property
	assert.Equal(t, sharedSecret1, sharedSecret2)
}

func TestComputeECDHSharedSecretInvalidPublicKey(t *testing.T) {
	// Generate a valid private key
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	assert.NoError(t, err)

	// Test with invalid public key length
	invalidPubKey := make([]byte, 64) // Should be 65 bytes
	_, err = ComputeECDHSharedSecret(privateKey, invalidPubKey)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "invalid peer public key format")

	// Test with invalid public key prefix
	invalidPrefixPubKey := make([]byte, 65)
	invalidPrefixPubKey[0] = 0x03 // Should be 0x04
	_, err = ComputeECDHSharedSecret(privateKey, invalidPrefixPubKey)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "invalid peer public key format")

	// Test with invalid point (not on curve)
	invalidPointPubKey := make([]byte, 65)
	invalidPointPubKey[0] = 0x04
	for i := 1; i < 65; i++ {
		invalidPointPubKey[i] = 0xFF // This point is not on P-256 curve
	}
	_, err = ComputeECDHSharedSecret(privateKey, invalidPointPubKey)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "failed to create ECDH public key")
}

func TestDeriveSessionKey(t *testing.T) {
	// Generate test data
	sharedSecret := make([]byte, 32)
	_, err := rand.Read(sharedSecret)
	assert.NoError(t, err)

	clientPubKey := make([]byte, 65)
	_, err = rand.Read(clientPubKey)
	assert.NoError(t, err)

	serverPubKey := make([]byte, 65)
	_, err = rand.Read(serverPubKey)
	assert.NoError(t, err)

	// Test session key derivation
	sessionKey, err := DeriveSessionKey(sharedSecret, clientPubKey, serverPubKey)
	assert.NoError(t, err)
	assert.NotNil(t, sessionKey)
	assert.Equal(t, 32, len(sessionKey)) // AES-256 key length

	// Test deterministic behavior
	sessionKey2, err := DeriveSessionKey(sharedSecret, clientPubKey, serverPubKey)
	assert.NoError(t, err)
	assert.Equal(t, sessionKey, sessionKey2)

	// Test different inputs produce different keys
	differentSharedSecret := make([]byte, 32)
	_, err = rand.Read(differentSharedSecret)
	assert.NoError(t, err)

	differentSessionKey, err := DeriveSessionKey(differentSharedSecret, clientPubKey, serverPubKey)
	assert.NoError(t, err)
	assert.NotEqual(t, sessionKey, differentSessionKey)
}

func TestGenerateKeyID(t *testing.T) {
	// Test key ID generation
	keyID1, err := GenerateKeyID()
	assert.NoError(t, err)
	assert.NotEmpty(t, keyID1)
	assert.Equal(t, 32, len(keyID1)) // 16 bytes hex encoded = 32 chars

	// Test uniqueness
	keyID2, err := GenerateKeyID()
	assert.NoError(t, err)
	assert.NotEqual(t, keyID1, keyID2)

	// Test format (should be hex)
	for _, char := range keyID1 {
		assert.True(t, (char >= '0' && char <= '9') || (char >= 'a' && char <= 'f'))
	}
}
