package cipherx_test

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/des"
	"errors"
	"testing"

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

// newBlock 按分组长度需要构造 AES（16 字节分组）或 DES（8 字节分组）分组密码。
func newBlock(useDES bool, key []byte) (cipher.Block, error) {
	if useDES {
		return des.NewCipher(key)
	}
	return aes.NewCipher(key)
}

func mustBlock(t *testing.T, useDES bool, key []byte) cipher.Block {
	t.Helper()
	b, err := newBlock(useDES, key)
	if err != nil {
		t.Fatalf("构造分组失败: %v", err)
	}
	return b
}

func TestModeMetadata(t *testing.T) {
	cases := []struct {
		mode     cipherx.Mode
		name     string
		usesIV   bool
		needsPad bool
	}{
		{cipherx.ECB, "ECB", false, true},
		{cipherx.CBC, "CBC", true, true},
		{cipherx.CFB, "CFB", true, false},
		{cipherx.OFB, "OFB", true, false},
		{cipherx.CTR, "CTR", true, false},
	}
	for _, c := range cases {
		if got := c.mode.String(); got != c.name {
			t.Errorf("Mode(%d).String() = %q, 期望 %q", int(c.mode), got, c.name)
		}
		if !c.mode.Valid() {
			t.Errorf("%s.Valid() = false, 期望 true", c.name)
		}
		if got := c.mode.UsesIV(); got != c.usesIV {
			t.Errorf("%s.UsesIV() = %v, 期望 %v", c.name, got, c.usesIV)
		}
		if got := c.mode.RequiresPadding(); got != c.needsPad {
			t.Errorf("%s.RequiresPadding() = %v, 期望 %v", c.name, got, c.needsPad)
		}
	}

	if cipherx.Mode(99).Valid() {
		t.Error("Mode(99).Valid() = true, 期望 false")
	}
	if got := cipherx.Mode(99).String(); got != "Mode(99)" {
		t.Errorf("Mode(99).String() = %q, 期望 %q", got, "Mode(99)")
	}
}

func TestPKCS7PadUnpadRoundTrip(t *testing.T) {
	for _, blockSize := range []int{8, 16} {
		for length := 0; length <= 2*blockSize+1; length++ {
			original := bytes.Repeat([]byte{0xAB}, length)
			padded, err := cipherx.PKCS7Pad(original, blockSize)
			if err != nil {
				t.Fatalf("PKCS7Pad(len=%d, bs=%d) 失败: %v", length, blockSize, err)
			}
			if len(padded)%blockSize != 0 {
				t.Fatalf("填充后长度 %d 不是 %d 的整数倍", len(padded), blockSize)
			}
			if len(padded) <= len(original) {
				if length%blockSize == 0 {
					// 恰好对齐时必须额外填充一整个分组
					t.Errorf("len=%d bs=%d 对齐时未补充完整分组", length, blockSize)
				}
			}
			// 入参不得被修改
			if !bytes.Equal(original, bytes.Repeat([]byte{0xAB}, length)) {
				t.Fatalf("PKCS7Pad 修改了入参")
			}

			unpadded, err := cipherx.PKCS7Unpad(padded, blockSize)
			if err != nil {
				t.Fatalf("PKCS7Unpad(len=%d, bs=%d) 失败: %v", length, blockSize, err)
			}
			if !bytes.Equal(unpadded, original) {
				t.Errorf("往返不一致: 得到 %x, 期望 %x", unpadded, original)
			}
		}
	}
}

func TestPKCS7PadInvalidBlockSize(t *testing.T) {
	for _, blockSize := range []int{-1, 0, 256, 1000} {
		if _, err := cipherx.PKCS7Pad([]byte("x"), blockSize); err == nil {
			t.Errorf("PKCS7Pad(bs=%d) 期望错误, 得到 nil", blockSize)
		} else if !errors.Is(err, cipherx.ErrInvalidBlockSize) {
			t.Errorf("PKCS7Pad(bs=%d) 错误类型不对: %v", blockSize, err)
		}
		if _, err := cipherx.PKCS7Unpad([]byte("12345678"), blockSize); err == nil {
			t.Errorf("PKCS7Unpad(bs=%d) 期望错误, 得到 nil", blockSize)
		}
	}
}

// TestPKCS7UnpadUniformError 验证所有填充失败原因都返回同一个错误，
// 以免调用方通过错误信息区分失败原因（padding oracle）。
func TestPKCS7UnpadUniformError(t *testing.T) {
	cases := []struct {
		name string
		data []byte
	}{
		{"空数据", nil},
		{"长度不是分组整数倍", []byte{1, 2, 3}},
		{"填充字节为 0", append(bytes.Repeat([]byte{0x00}, 15), 0x00)},
		{"填充长度超出分组", append(bytes.Repeat([]byte{0x00}, 15), 0x11)},
		{"填充字节不一致", append(bytes.Repeat([]byte{0x00}, 13), 0x03, 0x03, 0x04)},
		{"全零块", bytes.Repeat([]byte{0x00}, 16)},
	}
	for _, c := range cases {
		_, err := cipherx.PKCS7Unpad(c.data, 16)
		if err == nil {
			t.Errorf("%s: 期望错误, 得到 nil", c.name)
			continue
		}
		if !errors.Is(err, cipherx.ErrInvalidPadding) {
			t.Errorf("%s: 期望 ErrInvalidPadding, 得到 %v", c.name, err)
		}
		if err.Error() != cipherx.ErrInvalidPadding.Error() {
			t.Errorf("%s: 错误信息泄漏了失败原因: %q", c.name, err.Error())
		}
	}
}

// TestPKCS7UnpadReturnsCopy 验证去填充返回的切片不与入参共享底层数组。
func TestPKCS7UnpadReturnsCopy(t *testing.T) {
	data := append(bytes.Repeat([]byte{0x41}, 15), 0x01)
	out, err := cipherx.PKCS7Unpad(data, 16)
	if err != nil {
		t.Fatalf("PKCS7Unpad 失败: %v", err)
	}
	out[0] = 0xFF
	if data[0] != 0x41 {
		t.Error("PKCS7Unpad 返回的切片与入参共享底层数组")
	}
}

// TestAllModesRoundTrip 覆盖所有模式 × 两种分组长度 × 多种明文长度。
func TestAllModesRoundTrip(t *testing.T) {
	modes := []cipherx.Mode{cipherx.ECB, cipherx.CBC, cipherx.CFB, cipherx.OFB, cipherx.CTR}

	plaintexts := [][]byte{
		{},
		[]byte("a"),
		[]byte("0123456789abcde"),   // 15 字节
		[]byte("0123456789abcdef"),  // 16 字节
		[]byte("0123456789abcdefg"), // 17 字节
		bytes.Repeat([]byte("x"), 1000),
	}

	for _, useDES := range []bool{false, true} {
		var key []byte
		if useDES {
			key = bytes.Repeat([]byte{0x22}, 8)
		} else {
			key = bytes.Repeat([]byte{0x11}, 16)
		}
		algName := "AES"
		if useDES {
			algName = "DES"
		}

		for _, mode := range modes {
			for i, plaintext := range plaintexts {
				block, err := newBlock(useDES, key)
				if err != nil {
					t.Fatalf("构造分组失败: %v", err)
				}
				ciphertext, err := cipherx.Encrypt(block, mode, plaintext)
				if err != nil {
					t.Fatalf("%s-%s 明文#%d 加密失败: %v", algName, mode, i, err)
				}

				block2, err := newBlock(useDES, key)
				if err != nil {
					t.Fatalf("构造分组失败: %v", err)
				}
				decrypted, err := cipherx.Decrypt(block2, mode, ciphertext)
				if err != nil {
					t.Fatalf("%s-%s 明文#%d 解密失败: %v", algName, mode, i, err)
				}
				if !bytes.Equal(decrypted, plaintext) {
					t.Errorf("%s-%s 明文#%d 往返不一致: 得到 %q, 期望 %q",
						algName, mode, i, decrypted, plaintext)
				}

				// IV 前置格式：除 ECB 外，密文前缀必须等于 IV
				if mode.UsesIV() {
					bs := block.BlockSize()
					if len(ciphertext) < bs {
						t.Fatalf("%s-%s 密文短于分组长度", algName, mode)
					}
					iv := ciphertext[:bs]
					body := ciphertext[bs:]
					block3, _ := newBlock(useDES, key)
					decrypted2, err := cipherx.DecryptWithIV(block3, mode, body, iv)
					if err != nil {
						t.Fatalf("%s-%s WithIV 解密失败: %v", algName, mode, err)
					}
					if !bytes.Equal(decrypted2, plaintext) {
						t.Errorf("%s-%s WithIV 往返不一致", algName, mode)
					}
				}
			}
		}
	}
}

// TestMalformedCiphertextNeverPanics 是 P0-3 的回归测试。
//
// 旧实现在 ECB 解密时直接用 len(密文) 循环并按分组切片，密文长度不是分组
// 整数倍时会触发 "slice bounds out of range" panic，攻击者用一段畸形密文即可
// 让服务崩溃。现在必须返回 ErrNotBlockAligned。
func TestMalformedCiphertextNeverPanics(t *testing.T) {
	modes := []cipherx.Mode{cipherx.ECB, cipherx.CBC, cipherx.CFB, cipherx.OFB, cipherx.CTR}

	for _, useDES := range []bool{false, true} {
		key := bytes.Repeat([]byte{0x33}, 16)
		bs := 16
		if useDES {
			key = bytes.Repeat([]byte{0x33}, 8)
			bs = 8
		}
		block := mustBlock(t, useDES, key)

		for _, mode := range modes {
			// 长度 0 到 3 个分组 + 半个分组，全部枚举
			for n := 0; n <= 3*bs+bs/2; n++ {
				input := bytes.Repeat([]byte{0x7F}, n)
				// 不应 panic；返回什么都行，但必须安全返回
				_, _ = cipherx.Decrypt(block, mode, input)
				_, _ = cipherx.DecryptWithIV(block, mode, input, bytes.Repeat([]byte{0x01}, bs))
			}
		}
	}
}

// TestECBAndCBCRejectNonAligned 精确断言非对齐密文返回 ErrNotBlockAligned。
func TestECBAndCBCRejectNonAligned(t *testing.T) {
	block := mustBlock(t, false, bytes.Repeat([]byte{0x11}, 16))

	// 20 字节：不是 16 的整数倍，正是旧实现 panic 的输入
	nonAligned := bytes.Repeat([]byte{0xAB}, 20)

	if _, err := cipherx.Decrypt(block, cipherx.ECB, nonAligned); !errors.Is(err, cipherx.ErrNotBlockAligned) {
		t.Errorf("ECB 非对齐密文: 期望 ErrNotBlockAligned, 得到 %v", err)
	}

	// CBC 的 IV 前置格式：20 字节去掉 16 字节 IV 后剩 4 字节，同样非对齐
	if _, err := cipherx.Decrypt(block, cipherx.CBC, nonAligned); !errors.Is(err, cipherx.ErrNotBlockAligned) {
		t.Errorf("CBC 非对齐密文: 期望 ErrNotBlockAligned, 得到 %v", err)
	}
}

func TestEmptyCiphertext(t *testing.T) {
	block := mustBlock(t, false, bytes.Repeat([]byte{0x11}, 16))
	for _, mode := range []cipherx.Mode{cipherx.ECB, cipherx.CBC} {
		if _, err := cipherx.Decrypt(block, mode, nil); err == nil {
			t.Errorf("%s 空密文: 期望错误, 得到 nil", mode)
		}
	}
	// IV 模式：空密文不足以容纳 IV
	for _, mode := range []cipherx.Mode{cipherx.CFB, cipherx.OFB, cipherx.CTR} {
		if _, err := cipherx.Decrypt(block, mode, nil); !errors.Is(err, cipherx.ErrCiphertextTooShort) {
			t.Errorf("%s 空密文: 期望 ErrCiphertextTooShort, 得到 %v", mode, err)
		}
	}
}

func TestNilBlockAndInvalidMode(t *testing.T) {
	if _, err := cipherx.Encrypt(nil, cipherx.CBC, []byte("x")); !errors.Is(err, cipherx.ErrNilBlock) {
		t.Errorf("nil block: 期望 ErrNilBlock, 得到 %v", err)
	}
	if _, err := cipherx.Decrypt(nil, cipherx.CBC, []byte("x")); !errors.Is(err, cipherx.ErrNilBlock) {
		t.Errorf("nil block: 期望 ErrNilBlock, 得到 %v", err)
	}

	block := mustBlock(t, false, bytes.Repeat([]byte{0x11}, 16))
	if _, err := cipherx.Encrypt(block, cipherx.Mode(42), []byte("x")); !errors.Is(err, cipherx.ErrInvalidMode) {
		t.Errorf("非法模式: 期望 ErrInvalidMode, 得到 %v", err)
	}
	if _, err := cipherx.Decrypt(block, cipherx.Mode(-1), []byte("x")); !errors.Is(err, cipherx.ErrInvalidMode) {
		t.Errorf("非法模式: 期望 ErrInvalidMode, 得到 %v", err)
	}
}

func TestECBRejectsIV(t *testing.T) {
	block := mustBlock(t, false, bytes.Repeat([]byte{0x11}, 16))
	iv := bytes.Repeat([]byte{0x02}, 16)
	if _, err := cipherx.EncryptWithIV(block, cipherx.ECB, []byte("x"), iv); !errors.Is(err, cipherx.ErrInvalidMode) {
		t.Errorf("ECB WithIV: 期望 ErrInvalidMode, 得到 %v", err)
	}
	if _, err := cipherx.DecryptWithIV(block, cipherx.ECB, []byte("x"), iv); !errors.Is(err, cipherx.ErrInvalidMode) {
		t.Errorf("ECB WithIV 解密: 期望 ErrInvalidMode, 得到 %v", err)
	}
}

func TestInvalidIVSize(t *testing.T) {
	block := mustBlock(t, false, bytes.Repeat([]byte{0x11}, 16))
	for _, ivLen := range []int{0, 8, 15, 17, 32} {
		iv := bytes.Repeat([]byte{0x02}, ivLen)
		if _, err := cipherx.EncryptWithIV(block, cipherx.CBC, []byte("x"), iv); !errors.Is(err, cipherx.ErrInvalidIVSize) {
			t.Errorf("IV 长度 %d: 期望 ErrInvalidIVSize, 得到 %v", ivLen, err)
		}
	}
}

// TestRandomIVPerEncryption 验证相同明文两次加密得到不同密文（IV 随机）。
func TestRandomIVPerEncryption(t *testing.T) {
	block := mustBlock(t, false, bytes.Repeat([]byte{0x11}, 16))
	plaintext := []byte("same plaintext")

	seen := make(map[string]bool)
	for i := 0; i < 32; i++ {
		ciphertext, err := cipherx.Encrypt(block, cipherx.CBC, plaintext)
		if err != nil {
			t.Fatalf("加密失败: %v", err)
		}
		if seen[string(ciphertext)] {
			t.Fatal("相同明文产生了完全相同的密文，IV 可能未随机化")
		}
		seen[string(ciphertext)] = true
	}
}

// TestWrongKeyDoesNotPanic 验证错误密钥不会 panic（CBC 通常会因填充校验失败报错）。
func TestWrongKeyDoesNotPanic(t *testing.T) {
	right := mustBlock(t, false, bytes.Repeat([]byte{0x11}, 16))
	wrong := mustBlock(t, false, bytes.Repeat([]byte{0x99}, 16))

	ciphertext, err := cipherx.Encrypt(right, cipherx.CBC, []byte("secret data"))
	if err != nil {
		t.Fatalf("加密失败: %v", err)
	}
	// 允许返回错误（填充校验失败）或返回不同的明文，但绝不能 panic 或返回原文
	decrypted, err := cipherx.Decrypt(wrong, cipherx.CBC, ciphertext)
	if err == nil && bytes.Equal(decrypted, []byte("secret data")) {
		t.Error("错误密钥解出了正确明文")
	}
}

func TestHexHelpers(t *testing.T) {
	block := mustBlock(t, false, bytes.Repeat([]byte{0x11}, 16))
	plaintext := []byte("hex round trip 中文")

	for _, mode := range []cipherx.Mode{cipherx.ECB, cipherx.CBC, cipherx.CFB, cipherx.OFB, cipherx.CTR} {
		hexStr, err := cipherx.EncryptHex(block, mode, plaintext)
		if err != nil {
			t.Fatalf("%s EncryptHex 失败: %v", mode, err)
		}
		// 必须是大写、无换行
		for _, r := range hexStr {
			if (r < '0' || r > '9') && (r < 'A' || r > 'F') {
				t.Fatalf("%s EncryptHex 输出含非大写十六进制字符: %q", mode, hexStr)
			}
		}

		decrypted, err := cipherx.DecryptHex(block, mode, hexStr)
		if err != nil {
			t.Fatalf("%s DecryptHex 失败: %v", mode, err)
		}
		if !bytes.Equal(decrypted, plaintext) {
			t.Errorf("%s Hex 往返不一致", mode)
		}
	}

	if _, err := cipherx.DecryptHex(block, cipherx.CBC, "这不是十六进制"); err == nil {
		t.Error("非法十六进制: 期望错误, 得到 nil")
	}
	if _, err := cipherx.DecryptHex(block, cipherx.CBC, "ABC"); err == nil {
		t.Error("奇数长度十六进制: 期望错误, 得到 nil")
	}
}

func TestNewIV(t *testing.T) {
	iv, err := cipherx.NewIV(16)
	if err != nil {
		t.Fatalf("NewIV 失败: %v", err)
	}
	if len(iv) != 16 {
		t.Errorf("NewIV(16) 长度 = %d, 期望 16", len(iv))
	}
	if _, err := cipherx.NewIV(0); err == nil {
		t.Error("NewIV(0): 期望错误, 得到 nil")
	}
}