package desx_test

import (
	"bytes"
	"errors"
	"testing"

	"code.yun.ink/pkg/encryptx/cipherx"
	"code.yun.ink/pkg/encryptx/desx"
)

func TestGenerateKey(t *testing.T) {
	key, err := desx.GenerateKey()
	if err != nil {
		t.Fatalf("GenerateKey 失败: %v", err)
	}
	if len(key) != desx.KeySize {
		t.Fatalf("密钥长度 = %d, 期望 %d", len(key), desx.KeySize)
	}

	other, _ := desx.GenerateKey()
	if bytes.Equal(key, other) {
		t.Error("两次 GenerateKey 产生了相同密钥")
	}

	hexKey, err := desx.GenerateKeyHex()
	if err != nil {
		t.Fatalf("GenerateKeyHex 失败: %v", err)
	}
	if len(hexKey) != desx.KeySize*2 {
		t.Errorf("GenerateKeyHex 长度 = %d, 期望 %d", len(hexKey), desx.KeySize*2)
	}
}

func TestKeySizeValidation(t *testing.T) {
	plaintext := []byte("hello")
	for _, wrongLen := range []int{0, 7, 9, 16, 24} {
		key := bytes.Repeat([]byte{0x01}, wrongLen)
		if _, err := desx.EncryptCBC(plaintext, key); !errors.Is(err, cipherx.ErrInvalidKeySize) {
			t.Errorf("密钥长度 %d: 期望 ErrInvalidKeySize, 得到 %v", wrongLen, err)
		}
	}
}

func TestAllModesRoundTrip(t *testing.T) {
	key, err := desx.GenerateKey()
	if err != nil {
		t.Fatalf("GenerateKey 失败: %v", err)
	}

	encryptors := map[string]func([]byte, []byte) ([]byte, error){
		"ECB": desx.EncryptECB, "CBC": desx.EncryptCBC, "CFB": desx.EncryptCFB,
		"OFB": desx.EncryptOFB, "CTR": desx.EncryptCTR,
	}
	decryptors := map[string]func([]byte, []byte) ([]byte, error){
		"ECB": desx.DecryptECB, "CBC": desx.DecryptCBC, "CFB": desx.DecryptCFB,
		"OFB": desx.DecryptOFB, "CTR": desx.DecryptCTR,
	}
	plaintexts := [][]byte{
		{}, []byte("a"), []byte("1234567"), []byte("12345678"), []byte("123456789"),
		bytes.Repeat([]byte("DES"), 200),
	}

	for name, encrypt := range encryptors {
		for i, plaintext := range plaintexts {
			ciphertext, err := encrypt(plaintext, key)
			if err != nil {
				t.Fatalf("%s 明文#%d 加密失败: %v", name, i, err)
			}
			decrypted, err := decryptors[name](ciphertext, key)
			if err != nil {
				t.Fatalf("%s 明文#%d 解密失败: %v", name, i, err)
			}
			if !bytes.Equal(decrypted, plaintext) {
				t.Errorf("%s 明文#%d 往返不一致", name, i)
			}
		}
	}
}

func TestHexAPIRoundTrip(t *testing.T) {
	key, _ := desx.GenerateKey()
	plaintext := []byte("十六进制往返")

	cases := []struct {
		name    string
		encrypt func([]byte, []byte) (string, error)
		decrypt func(string, []byte) ([]byte, error)
	}{
		{"ECB", desx.EncryptECBHex, desx.DecryptECBHex},
		{"CBC", desx.EncryptCBCHex, desx.DecryptCBCHex},
		{"CFB", desx.EncryptCFBHex, desx.DecryptCFBHex},
		{"OFB", desx.EncryptOFBHex, desx.DecryptOFBHex},
		{"CTR", desx.EncryptCTRHex, desx.DecryptCTRHex},
	}
	for _, c := range cases {
		hexStr, err := c.encrypt(plaintext, key)
		if err != nil {
			t.Fatalf("%s Hex 加密失败: %v", c.name, err)
		}
		if bytes.ContainsAny([]byte(hexStr), "\n\r \t") {
			t.Errorf("%s Hex 输出含空白字符", c.name)
		}
		decrypted, err := c.decrypt(hexStr, key)
		if err != nil {
			t.Fatalf("%s Hex 解密失败: %v", c.name, err)
		}
		if !bytes.Equal(decrypted, plaintext) {
			t.Errorf("%s Hex 往返不一致", c.name)
		}
	}
}

func TestWithIVRoundTrip(t *testing.T) {
	key, _ := desx.GenerateKey()
	iv := bytes.Repeat([]byte{0x5A}, desx.BlockSize)
	plaintext := []byte("fixed iv")

	cases := []struct {
		name    string
		encrypt func([]byte, []byte, []byte) ([]byte, error)
		decrypt func([]byte, []byte, []byte) ([]byte, error)
	}{
		{"CBC", desx.EncryptCBCWithIV, desx.DecryptCBCWithIV},
		{"CFB", desx.EncryptCFBWithIV, desx.DecryptCFBWithIV},
		{"OFB", desx.EncryptOFBWithIV, desx.DecryptOFBWithIV},
		{"CTR", desx.EncryptCTRWithIV, desx.DecryptCTRWithIV},
	}
	for _, c := range cases {
		ciphertext, err := c.encrypt(plaintext, key, iv)
		if err != nil {
			t.Fatalf("%s WithIV 加密失败: %v", c.name, err)
		}
		decrypted, err := c.decrypt(ciphertext, key, iv)
		if err != nil {
			t.Fatalf("%s WithIV 解密失败: %v", c.name, err)
		}
		if !bytes.Equal(decrypted, plaintext) {
			t.Errorf("%s WithIV 往返不一致", c.name)
		}
	}
}

// TestMalformedCiphertextNeverPanics 是 P0-3 在 DES 上的回归测试。
func TestMalformedCiphertextNeverPanics(t *testing.T) {
	key, _ := desx.GenerateKey()
	decryptors := map[string]func([]byte, []byte) ([]byte, error){
		"ECB": desx.DecryptECB, "CBC": desx.DecryptCBC, "CFB": desx.DecryptCFB,
		"OFB": desx.DecryptOFB, "CTR": desx.DecryptCTR,
	}
	for _, decrypt := range decryptors {
		for n := 0; n <= 24; n++ {
			_, _ = decrypt(bytes.Repeat([]byte{0x7F}, n), key)
		}
	}

	nonAligned := bytes.Repeat([]byte{0xAB}, 12)
	if _, err := desx.DecryptECB(nonAligned, key); !errors.Is(err, cipherx.ErrNotBlockAligned) {
		t.Errorf("ECB 非对齐: 期望 ErrNotBlockAligned, 得到 %v", err)
	}
	if _, err := desx.DecryptCBC(nonAligned, key); !errors.Is(err, cipherx.ErrNotBlockAligned) {
		t.Errorf("CBC 非对齐: 期望 ErrNotBlockAligned, 得到 %v", err)
	}
}

func TestECBIsDeterministic(t *testing.T) {
	key, _ := desx.GenerateKey()
	plaintext := bytes.Repeat([]byte("12345678"), 3)

	a, _ := desx.EncryptECB(plaintext, key)
	b, _ := desx.EncryptECB(plaintext, key)
	if !bytes.Equal(a, b) {
		t.Error("ECB 应当是确定性的")
	}

	c, _ := desx.EncryptCBC(plaintext, key)
	d, _ := desx.EncryptCBC(plaintext, key)
	if bytes.Equal(c, d) {
		t.Error("CBC 使用了随机 IV，两次结果不应相同")
	}
}