// 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 tdx

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"testing"
)

func TestParseECDSAPublicKey(t *testing.T) {
	// Generate a test key pair
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("Failed to generate test key: %v", err)
	}

	// Get the public key bytes in uncompressed format
	publicKeyBytes := elliptic.Marshal(elliptic.P256(), privateKey.PublicKey.X, privateKey.PublicKey.Y)

	// Remove the 0x04 prefix to get the raw 64 bytes
	rawKey := publicKeyBytes[1:]
	if len(rawKey) != 64 {
		t.Fatalf("Expected 64 bytes, got %d", len(rawKey))
	}

	// Test parsing
	parsedKey, err := ParseECDSAPublicKey(rawKey)
	if err != nil {
		t.Fatalf("ParseECDSAPublicKey failed: %v", err)
	}

	// Verify the parsed key matches the original
	if parsedKey.X.Cmp(privateKey.PublicKey.X) != 0 {
		t.Errorf("X coordinate mismatch: got %v, want %v", parsedKey.X, privateKey.PublicKey.X)
	}
	if parsedKey.Y.Cmp(privateKey.PublicKey.Y) != 0 {
		t.Errorf("Y coordinate mismatch: got %v, want %v", parsedKey.Y, privateKey.PublicKey.Y)
	}
}

func TestParseECDSAPublicKeyInvalidLength(t *testing.T) {
	// Test with invalid length
	_, err := ParseECDSAPublicKey([]byte("invalid"))
	if err == nil {
		t.Error("Expected error for invalid key length, but got nil")
	}
}

func TestVerifyECDSASignature(t *testing.T) {
	// Generate a test key pair
	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("Failed to generate test key: %v", err)
	}

	// Create a test message
	message := []byte("test message")

	// Sign the message
	hash := sha256.Sum256(message)
	r, s, err := ecdsa.Sign(rand.Reader, privateKey, hash[:])
	if err != nil {
		t.Fatalf("Failed to sign message: %v", err)
	}

	// Create signature in P1363 format (r || s, each 32 bytes)
	signature := make([]byte, 64)
	r.FillBytes(signature[:32])
	s.FillBytes(signature[32:])

	// Test verification with correct key
	err = VerifyECDSASignature(&privateKey.PublicKey, message, signature)
	if err != nil {
		t.Errorf("Verification failed with correct key: %v", err)
	}

	// Test verification with wrong key (generate a new key)
	wrongKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("Failed to generate wrong key: %v", err)
	}

	err = VerifyECDSASignature(&wrongKey.PublicKey, message, signature)
	if err == nil {
		t.Error("Expected verification to fail with wrong key, but it succeeded")
	}
}
