package asymm

import (
	"bytes"
	"testing"
)

func TestSM2SignVerifyAndEncryptDecrypt(t *testing.T) {
	priv, pub, err := GenerateSM2Key()
	if err != nil {
		t.Fatalf("GenerateSM2Key failed: %v", err)
	}

	msg := []byte("sm2-message")
	uid := []byte("user123")
	sig, err := SM2Sign(priv, msg, uid)
	if err != nil {
		t.Fatalf("SM2Sign failed: %v", err)
	}
	if !SM2Verify(pub, msg, sig, uid) {
		t.Fatalf("SM2Verify failed")
	}

	cipher, err := SM2EncryptASN1(pub, msg)
	if err != nil {
		t.Fatalf("SM2EncryptASN1 failed: %v", err)
	}
	plain, err := SM2DecryptASN1(priv, cipher)
	if err != nil {
		t.Fatalf("SM2DecryptASN1 failed: %v", err)
	}
	if !bytes.Equal(plain, msg) {
		t.Fatalf("SM2 decrypt mismatch")
	}
}

func TestSM2PEMEncodeDecode(t *testing.T) {
	priv, pub, err := GenerateSM2Key()
	if err != nil {
		t.Fatalf("GenerateSM2Key failed: %v", err)
	}

	privPEM, err := EncodeSM2PrivateKey(priv, "pwd")
	if err != nil {
		t.Fatalf("EncodeSM2PrivateKey failed: %v", err)
	}
	pubPEM, err := EncodeSM2PublicKey(pub)
	if err != nil {
		t.Fatalf("EncodeSM2PublicKey failed: %v", err)
	}

	decodedPriv, err := DecodeSM2PrivateKey(privPEM, "pwd")
	if err != nil {
		t.Fatalf("DecodeSM2PrivateKey failed: %v", err)
	}
	decodedPub, err := DecodeSM2PublicKey(pubPEM)
	if err != nil {
		t.Fatalf("DecodeSM2PublicKey failed: %v", err)
	}

	msg := []byte("sm2-pem")
	sig, err := SM2Sign(decodedPriv, msg, nil)
	if err != nil {
		t.Fatalf("SM2Sign failed: %v", err)
	}
	if !SM2Verify(decodedPub, msg, sig, nil) {
		t.Fatalf("SM2 verify with decoded keys failed")
	}
}

func TestSM9SignVerifyAndEncryptDecrypt(t *testing.T) {
	signMasterPriv, signMasterPub, err := GenerateSM9SignMasterKey()
	if err != nil {
		t.Fatalf("GenerateSM9SignMasterKey failed: %v", err)
	}
	encryptMasterPriv, encryptMasterPub, err := GenerateSM9EncryptMasterKey()
	if err != nil {
		t.Fatalf("GenerateSM9EncryptMasterKey failed: %v", err)
	}

	uid := []byte("alice@example.com")
	signUserKey, err := GenerateSM9SignUserKey(signMasterPriv, uid, 0)
	if err != nil {
		t.Fatalf("GenerateSM9SignUserKey failed: %v", err)
	}
	encryptUserKey, err := GenerateSM9EncryptUserKey(encryptMasterPriv, uid, 0)
	if err != nil {
		t.Fatalf("GenerateSM9EncryptUserKey failed: %v", err)
	}

	msg := []byte("sm9-message")
	sig, err := SM9SignASN1(signUserKey, msg)
	if err != nil {
		t.Fatalf("SM9SignASN1 failed: %v", err)
	}
	if !SM9VerifyASN1(signMasterPub, uid, 0, msg, sig) {
		t.Fatalf("SM9VerifyASN1 failed")
	}

	cipher, err := SM9Encrypt(encryptMasterPub, uid, 0, msg)
	if err != nil {
		t.Fatalf("SM9Encrypt failed: %v", err)
	}
	plain, err := SM9Decrypt(encryptUserKey, uid, cipher)
	if err != nil {
		t.Fatalf("SM9Decrypt failed: %v", err)
	}
	if !bytes.Equal(plain, msg) {
		t.Fatalf("SM9 decrypt mismatch")
	}
}

func TestSM9PEMEncodeDecode(t *testing.T) {
	signMasterPriv, _, err := GenerateSM9SignMasterKey()
	if err != nil {
		t.Fatalf("GenerateSM9SignMasterKey failed: %v", err)
	}
	encryptMasterPriv, _, err := GenerateSM9EncryptMasterKey()
	if err != nil {
		t.Fatalf("GenerateSM9EncryptMasterKey failed: %v", err)
	}

	signPrivPEM, err := EncodeSM9SignMasterPrivateKey(signMasterPriv)
	if err != nil {
		t.Fatalf("EncodeSM9SignMasterPrivateKey failed: %v", err)
	}
	encPrivPEM, err := EncodeSM9EncryptMasterPrivateKey(encryptMasterPriv)
	if err != nil {
		t.Fatalf("EncodeSM9EncryptMasterPrivateKey failed: %v", err)
	}

	if _, err := DecodeSM9SignMasterPrivateKey(signPrivPEM); err != nil {
		t.Fatalf("DecodeSM9SignMasterPrivateKey failed: %v", err)
	}
	if _, err := DecodeSM9EncryptMasterPrivateKey(encPrivPEM); err != nil {
		t.Fatalf("DecodeSM9EncryptMasterPrivateKey failed: %v", err)
	}
}