package starcrypto

import (
	"bytes"
	"testing"
)

func TestRootSymmetricWrappers(t *testing.T) {
	aesKey := []byte("0123456789abcdef")
	aesIV := []byte("abcdef9876543210")
	plain := []byte("root-wrapper-aes")

	aesEnc, err := EncryptAesCBC(plain, aesKey, aesIV, "")
	if err != nil {
		t.Fatalf("EncryptAesCBC failed: %v", err)
	}
	aesDec, err := DecryptAesCBC(aesEnc, aesKey, aesIV, "")
	if err != nil {
		t.Fatalf("DecryptAesCBC failed: %v", err)
	}
	if !bytes.Equal(aesDec, plain) {
		t.Fatalf("aes wrapper mismatch")
	}

	aesCFBEnc, err := EncryptAes(plain, aesKey, aesIV, MODECFB, "")
	if err != nil {
		t.Fatalf("EncryptAes failed: %v", err)
	}
	aesCFBDec, err := DecryptAes(aesCFBEnc, aesKey, aesIV, MODECFB, "")
	if err != nil {
		t.Fatalf("DecryptAes failed: %v", err)
	}
	if !bytes.Equal(aesCFBDec, plain) {
		t.Fatalf("aes generic wrapper mismatch")
	}

	aesStreamEnc := &bytes.Buffer{}
	if err := EncryptAesCBCStream(aesStreamEnc, bytes.NewReader(plain), aesKey, aesIV, ""); err != nil {
		t.Fatalf("EncryptAesCBCStream failed: %v", err)
	}
	aesStreamDec := &bytes.Buffer{}
	if err := DecryptAesCBCStream(aesStreamDec, bytes.NewReader(aesStreamEnc.Bytes()), aesKey, aesIV, ""); err != nil {
		t.Fatalf("DecryptAesCBCStream failed: %v", err)
	}
	if !bytes.Equal(aesStreamDec.Bytes(), plain) {
		t.Fatalf("aes stream wrapper mismatch")
	}

	sm4Enc, err := EncryptSM4CBC(plain, aesKey, aesIV, "")
	if err != nil {
		t.Fatalf("EncryptSM4CBC failed: %v", err)
	}
	sm4Dec, err := DecryptSM4CBC(sm4Enc, aesKey, aesIV, "")
	if err != nil {
		t.Fatalf("DecryptSM4CBC failed: %v", err)
	}
	if !bytes.Equal(sm4Dec, plain) {
		t.Fatalf("sm4 wrapper mismatch")
	}

	sm4StreamEnc := &bytes.Buffer{}
	if err := EncryptSM4CBCStream(sm4StreamEnc, bytes.NewReader(plain), aesKey, aesIV, ""); err != nil {
		t.Fatalf("EncryptSM4CBCStream failed: %v", err)
	}
	sm4StreamDec := &bytes.Buffer{}
	if err := DecryptSM4CBCStream(sm4StreamDec, bytes.NewReader(sm4StreamEnc.Bytes()), aesKey, aesIV, ""); err != nil {
		t.Fatalf("DecryptSM4CBCStream failed: %v", err)
	}
	if !bytes.Equal(sm4StreamDec.Bytes(), plain) {
		t.Fatalf("sm4 stream wrapper mismatch")
	}

	desKey := []byte("12345678")
	desIV := []byte("abcdefgh")
	desEnc, err := EncryptDESCBC(plain, desKey, desIV, "")
	if err != nil {
		t.Fatalf("EncryptDESCBC failed: %v", err)
	}
	desDec, err := DecryptDESCBC(desEnc, desKey, desIV, "")
	if err != nil {
		t.Fatalf("DecryptDESCBC failed: %v", err)
	}
	if !bytes.Equal(desDec, plain) {
		t.Fatalf("des wrapper mismatch")
	}

	desStreamEnc := &bytes.Buffer{}
	if err := EncryptDESCBCStream(desStreamEnc, bytes.NewReader(plain), desKey, desIV, ""); err != nil {
		t.Fatalf("EncryptDESCBCStream failed: %v", err)
	}
	desStreamDec := &bytes.Buffer{}
	if err := DecryptDESCBCStream(desStreamDec, bytes.NewReader(desStreamEnc.Bytes()), desKey, desIV, ""); err != nil {
		t.Fatalf("DecryptDESCBCStream failed: %v", err)
	}
	if !bytes.Equal(desStreamDec.Bytes(), plain) {
		t.Fatalf("des stream wrapper mismatch")
	}
}

func TestRootSM2AndSM9Wrappers(t *testing.T) {
	sm2Priv, sm2Pub, err := GenerateSM2Key()
	if err != nil {
		t.Fatalf("GenerateSM2Key failed: %v", err)
	}
	msg := []byte("root-sm")
	sig, err := SM2Sign(sm2Priv, msg, nil)
	if err != nil {
		t.Fatalf("SM2Sign failed: %v", err)
	}
	if !SM2Verify(sm2Pub, msg, sig, nil) {
		t.Fatalf("SM2Verify failed")
	}

	signMaster, signPub, err := GenerateSM9SignMasterKey()
	if err != nil {
		t.Fatalf("GenerateSM9SignMasterKey failed: %v", err)
	}
	encryptMaster, encryptPub, err := GenerateSM9EncryptMasterKey()
	if err != nil {
		t.Fatalf("GenerateSM9EncryptMasterKey failed: %v", err)
	}
	uid := []byte("root@example.com")
	signUser, err := GenerateSM9SignUserKey(signMaster, uid, 0)
	if err != nil {
		t.Fatalf("GenerateSM9SignUserKey failed: %v", err)
	}
	encryptUser, err := GenerateSM9EncryptUserKey(encryptMaster, uid, 0)
	if err != nil {
		t.Fatalf("GenerateSM9EncryptUserKey failed: %v", err)
	}

	sm9Sig, err := SM9SignASN1(signUser, msg)
	if err != nil {
		t.Fatalf("SM9SignASN1 failed: %v", err)
	}
	if !SM9VerifyASN1(signPub, uid, 0, msg, sm9Sig) {
		t.Fatalf("SM9VerifyASN1 failed")
	}

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

func TestRootChaChaAndSM4DerivedWrappers(t *testing.T) {
	key := []byte("0123456789abcdef0123456789abcdef")
	nonce := []byte("123456789012")
	aad := []byte("aad")
	plain := []byte("root-chacha-wrapper")

	chachaEnc, err := EncryptChaCha20(plain, key, nonce)
	if err != nil {
		t.Fatalf("EncryptChaCha20 failed: %v", err)
	}
	chachaDec, err := DecryptChaCha20(chachaEnc, key, nonce)
	if err != nil {
		t.Fatalf("DecryptChaCha20 failed: %v", err)
	}
	if !bytes.Equal(chachaDec, plain) {
		t.Fatalf("chacha wrapper mismatch")
	}

	polyEnc, err := EncryptChaCha20Poly1305(plain, key, nonce, aad)
	if err != nil {
		t.Fatalf("EncryptChaCha20Poly1305 failed: %v", err)
	}
	polyDec, err := DecryptChaCha20Poly1305(polyEnc, key, nonce, aad)
	if err != nil {
		t.Fatalf("DecryptChaCha20Poly1305 failed: %v", err)
	}
	if !bytes.Equal(polyDec, plain) {
		t.Fatalf("chacha20-poly1305 wrapper mismatch")
	}

	sm4Key := []byte("0123456789abcdef")
	sm4IV := []byte("abcdef9876543210")
	sm4Plain := []byte("root-sm4-derived")

	ecbEnc, err := EncryptSM4ECB(sm4Plain, sm4Key, "")
	if err != nil {
		t.Fatalf("EncryptSM4ECB failed: %v", err)
	}
	ecbDec, err := DecryptSM4ECB(ecbEnc, sm4Key, "")
	if err != nil {
		t.Fatalf("DecryptSM4ECB failed: %v", err)
	}
	if !bytes.Equal(ecbDec, sm4Plain) {
		t.Fatalf("sm4 ecb wrapper mismatch")
	}

	ofbEnc, err := EncryptSM4OFB(sm4Plain, sm4Key, sm4IV)
	if err != nil {
		t.Fatalf("EncryptSM4OFB failed: %v", err)
	}
	ofbDec, err := DecryptSM4OFB(ofbEnc, sm4Key, sm4IV)
	if err != nil {
		t.Fatalf("DecryptSM4OFB failed: %v", err)
	}
	if !bytes.Equal(ofbDec, sm4Plain) {
		t.Fatalf("sm4 ofb wrapper mismatch")
	}

	ctrEnc, err := EncryptSM4CTR(sm4Plain, sm4Key, sm4IV)
	if err != nil {
		t.Fatalf("EncryptSM4CTR failed: %v", err)
	}
	ctrDec, err := DecryptSM4CTR(ctrEnc, sm4Key, sm4IV)
	if err != nil {
		t.Fatalf("DecryptSM4CTR failed: %v", err)
	}
	if !bytes.Equal(ctrDec, sm4Plain) {
		t.Fatalf("sm4 ctr wrapper mismatch")
	}

	ecbStreamEnc := &bytes.Buffer{}
	if err := EncryptSM4ECBStream(ecbStreamEnc, bytes.NewReader(sm4Plain), sm4Key, ""); err != nil {
		t.Fatalf("EncryptSM4ECBStream failed: %v", err)
	}
	ecbStreamDec := &bytes.Buffer{}
	if err := DecryptSM4ECBStream(ecbStreamDec, bytes.NewReader(ecbStreamEnc.Bytes()), sm4Key, ""); err != nil {
		t.Fatalf("DecryptSM4ECBStream failed: %v", err)
	}
	if !bytes.Equal(ecbStreamDec.Bytes(), sm4Plain) {
		t.Fatalf("sm4 ecb stream wrapper mismatch")
	}
}

func TestRootOptionsAndGCMWrappers(t *testing.T) {
	aesKey := []byte("0123456789abcdef")
	sm4Key := []byte("0123456789abcdef")
	nonce := []byte("123456789012")
	plain := []byte("root-options-gcm")

	aesEnc, err := EncryptAesGCM(plain, aesKey, nonce, []byte("aad"))
	if err != nil {
		t.Fatalf("EncryptAesGCM failed: %v", err)
	}
	aesDec, err := DecryptAesGCM(aesEnc, aesKey, nonce, []byte("aad"))
	if err != nil {
		t.Fatalf("DecryptAesGCM failed: %v", err)
	}
	if !bytes.Equal(aesDec, plain) {
		t.Fatalf("aes gcm wrapper mismatch")
	}

	sm4Enc, err := EncryptSM4WithOptions(plain, sm4Key, &CipherOptions{Nonce: nonce})
	if err != nil {
		t.Fatalf("EncryptSM4WithOptions failed: %v", err)
	}
	sm4Dec, err := DecryptSM4WithOptions(sm4Enc, sm4Key, &CipherOptions{Nonce: nonce})
	if err != nil {
		t.Fatalf("DecryptSM4WithOptions failed: %v", err)
	}
	if !bytes.Equal(sm4Dec, plain) {
		t.Fatalf("sm4 options wrapper mismatch")
	}
}