// 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/aes"
	"crypto/cipher"
	"crypto/ecdh"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/asn1"
	"encoding/pem"
	"fmt"
	"io"
	"math/big"

	"golang.org/x/crypto/hkdf"
)

// ConvertRSAPubKeyToPEM converts an RSA public key to PEM format
func ConvertRSAPubKeyToPEM(pubKey *rsa.PublicKey) (string, error) {
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		return "", fmt.Errorf("failed to marshal public key: %w", err)
	}

	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubKeyBytes,
	})

	return string(pubKeyPEM), nil
}

// ConvertRSAPubKeyToPEMBytes converts an RSA public key to PEM format and returns bytes
func ConvertRSAPubKeyToPEMBytes(pubKey *rsa.PublicKey) ([]byte, error) {
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal public key: %w", err)
	}

	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubKeyBytes,
	})

	return pubKeyPEM, nil
}

// ParseRSAPubKeyFromPEM parses an RSA public key from PEM format
func ParseRSAPubKeyFromPEM(pubKeyPEM string) (*rsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(pubKeyPEM))
	if block == nil {
		return nil, fmt.Errorf("failed to decode PEM block")
	}

	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse public key: %w", err)
	}

	rsaPubKey, ok := pubKey.(*rsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("parsed key is not an RSA public key")
	}

	return rsaPubKey, nil
}

// ConvertECDSAPubKeyToPEM converts an ECDSA public key to PEM format and returns string
func ConvertECDSAPubKeyToPEM(pubKey *ecdsa.PublicKey) (string, error) {
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		return "", fmt.Errorf("failed to marshal ECDSA public key: %w", err)
	}

	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubKeyBytes,
	})

	return string(pubKeyPEM), nil
}

// ConvertECDSAPubKeyToPEMBytes converts an ECDSA public key to PEM format and returns bytes
func ConvertECDSAPubKeyToPEMBytes(pubKey *ecdsa.PublicKey) ([]byte, error) {
	pubKeyBytes, err := x509.MarshalPKIXPublicKey(pubKey)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal ECDSA public key: %w", err)
	}

	pubKeyPEM := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubKeyBytes,
	})

	return pubKeyPEM, nil
}

// ParseECDSAPubKeyFromPEM parses an ECDSA public key from PEM format
func ParseECDSAPubKeyFromPEM(pubKeyPEM string) (*ecdsa.PublicKey, error) {
	block, _ := pem.Decode([]byte(pubKeyPEM))
	if block == nil {
		return nil, fmt.Errorf("failed to decode PEM block")
	}

	pubKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("failed to parse public key: %w", err)
	}

	ecdsaPubKey, ok := pubKey.(*ecdsa.PublicKey)
	if !ok {
		return nil, fmt.Errorf("parsed key is not an ECDSA public key")
	}

	return ecdsaPubKey, nil
}

// MarshalECDHPublicKey marshals an ECDSA public key to uncompressed format for ECDH
// This replaces the deprecated elliptic.Marshal function
func MarshalECDHPublicKey(pubKey *ecdsa.PublicKey) []byte {
	// Uncompressed format: 0x04 || X || Y (65 bytes for P-256)
	keyBytes := make([]byte, 65)
	keyBytes[0] = 0x04 // Uncompressed point indicator

	// Copy X coordinate (32 bytes)
	xBytes := pubKey.X.Bytes()
	copy(keyBytes[33-len(xBytes):33], xBytes)

	// Copy Y coordinate (32 bytes)
	yBytes := pubKey.Y.Bytes()
	copy(keyBytes[65-len(yBytes):65], yBytes)

	return keyBytes
}

// EncryptWithAESGCM encrypts data using AES-GCM with a 32-byte session key
func EncryptWithAESGCM(plaintext []byte, sessionKey []byte) ([]byte, error) {
	// Validate session key length (should be 32 bytes for AES-256)
	if len(sessionKey) != 32 {
		return nil, fmt.Errorf("invalid session key length: expected 32 bytes, got %d", len(sessionKey))
	}

	// Create cipher
	block, err := aes.NewCipher(sessionKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create AES cipher: %w", err)
	}

	// Create GCM
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	// Generate nonce
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, fmt.Errorf("failed to generate nonce: %w", err)
	}

	// Encrypt the data
	ciphertext := gcm.Seal(nonce, nonce, plaintext, nil)

	return ciphertext, nil
}

// DecryptWithAESGCM decrypts data using AES-GCM with a 32-byte session key
func DecryptWithAESGCM(ciphertext []byte, sessionKey []byte) ([]byte, error) {
	// Validate session key length (should be 32 bytes for AES-256)
	if len(sessionKey) != 32 {
		return nil, fmt.Errorf("invalid session key length: expected 32 bytes, got %d", len(sessionKey))
	}

	// Create cipher
	block, err := aes.NewCipher(sessionKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create AES cipher: %w", err)
	}

	// Create GCM
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("failed to create GCM: %w", err)
	}

	// Validate ciphertext length
	nonceSize := gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		return nil, fmt.Errorf("ciphertext too short")
	}

	// Extract nonce and ciphertext
	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]

	// Decrypt the data
	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to decrypt data: %w", err)
	}

	return plaintext, nil
}

// ComputeECDHSharedSecret computes the ECDH shared secret using modern crypto/ecdh package
func ComputeECDHSharedSecret(privateKey *ecdsa.PrivateKey, peerPublicKey []byte) ([]byte, error) {
	// Validate peer public key format (uncompressed format: 0x04 + 32 bytes X + 32 bytes Y)
	if len(peerPublicKey) != 65 || peerPublicKey[0] != 0x04 {
		return nil, fmt.Errorf("invalid peer public key format: expected 65 bytes with 0x04 prefix")
	}

	// Convert ECDSA private key to ECDH private key
	ecdhPrivateKey, err := privateKey.ECDH()
	if err != nil {
		return nil, fmt.Errorf("failed to convert ECDSA private key to ECDH: %w", err)
	}

	// Create ECDH public key from peer's public key bytes
	// The crypto/ecdh package expects the raw uncompressed format
	ecdhPublicKey, err := ecdh.P256().NewPublicKey(peerPublicKey)
	if err != nil {
		return nil, fmt.Errorf("failed to create ECDH public key from peer key: %w", err)
	}

	// Compute the shared secret using ECDH
	sharedSecret, err := ecdhPrivateKey.ECDH(ecdhPublicKey)
	if err != nil {
		return nil, fmt.Errorf("failed to compute ECDH shared secret: %w", err)
	}

	return sharedSecret, nil
}

// DeriveSessionKey derives a session key using HKDF-SHA256
func DeriveSessionKey(sharedSecret, clientPublicKey, serverPublicKey []byte) ([]byte, error) {
	// Use client and server public keys as info parameter
	info := append(clientPublicKey, serverPublicKey...)

	// Use empty salt (HKDF will use zero salt)
	salt := make([]byte, 0)

	// Create HKDF reader
	hkdfReader := hkdf.New(sha256.New, sharedSecret, salt, info)

	// Derive 32-byte session key
	sessionKey := make([]byte, 32)
	if _, err := io.ReadFull(hkdfReader, sessionKey); err != nil {
		return nil, fmt.Errorf("failed to derive session key: %w", err)
	}

	return sessionKey, nil
}

// GenerateKeyID generates a unique key ID
func GenerateKeyID() (string, error) {
	bytes := make([]byte, 16)
	if _, err := rand.Read(bytes); err != nil {
		return "", fmt.Errorf("failed to generate random bytes: %w", err)
	}

	return fmt.Sprintf("%x", bytes), nil
}

// ecdsaSignature represents an ECDSA signature in ASN.1 DER format
type ecdsaSignature struct {
	R, S *big.Int
}

// SignECDSAPublicKey signs an ECDSA public key with a signing private key
// Returns the signature in ASN.1 DER format
func SignECDSAPublicKey(pubKeyToSign *ecdsa.PublicKey, signingPrivKey *ecdsa.PrivateKey) ([]byte, error) {
	// Convert public key to PEM format
	pubKeyPEM, err := ConvertECDSAPubKeyToPEMBytes(pubKeyToSign)
	if err != nil {
		return nil, fmt.Errorf("failed to convert public key to PEM: %w", err)
	}

	// Compute SHA256 hash of the PEM-encoded public key
	hash := sha256.Sum256(pubKeyPEM)

	// Sign the hash using ECDSA
	r, s, err := ecdsa.Sign(rand.Reader, signingPrivKey, hash[:])
	if err != nil {
		return nil, fmt.Errorf("failed to sign data: %w", err)
	}

	// Encode signature in ASN.1 DER format
	signature, err := asn1.Marshal(ecdsaSignature{R: r, S: s})
	if err != nil {
		return nil, fmt.Errorf("failed to marshal signature: %w", err)
	}

	return signature, nil
}

// VerifyECDSAPublicKeySignature verifies the signature of an ECDSA public key
func VerifyECDSAPublicKeySignature(pubKey *ecdsa.PublicKey, signatureDER []byte, signingPubKey *ecdsa.PublicKey) error {
	// Parse the DER-encoded signature
	var sig ecdsaSignature
	if _, err := asn1.Unmarshal(signatureDER, &sig); err != nil {
		return fmt.Errorf("failed to unmarshal signature: %w", err)
	}

	// Convert public key to PEM format
	pubKeyPEM, err := ConvertECDSAPubKeyToPEMBytes(pubKey)
	if err != nil {
		return fmt.Errorf("failed to convert public key to PEM: %w", err)
	}

	// Compute SHA256 hash of the PEM-encoded public key
	hash := sha256.Sum256(pubKeyPEM)

	// Verify the signature
	if !ecdsa.Verify(signingPubKey, hash[:], sig.R, sig.S) {
		return fmt.Errorf("signature verification failed")
	}

	return nil
}
