package pbkdf2x_test

import (
	"bytes"
	"crypto/sha1"
	"crypto/sha512"
	"encoding/hex"
	"math"
	"testing"

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

// mustHexDecode 把测试向量里的十六进制串还原为字节。
func mustHexDecode(t *testing.T, s string) []byte {
	t.Helper()
	b, err := hex.DecodeString(s)
	if err != nil {
		t.Fatalf("测试向量不是合法十六进制: %v", err)
	}
	return b
}

// TestDeriveKeyKnownVectors 固定本包的已知答案（KAT）。
// 期望值由 Python 3.13 的 hashlib.pbkdf2_hmac（独立实现）生成，例如：
//
//	hashlib.pbkdf2_hmac('sha256', b'password', b'saltsaltsalt', 10000, 32).hex()
func TestDeriveKeyKnownVectors(t *testing.T) {
	cases := []struct {
		name       string
		password   []byte
		saltHex    string
		iterations int
		keyLength  int
		want       string
	}{
		{
			name:       "sha256-10000-32",
			password:   []byte("password"),
			saltHex:    hex.EncodeToString([]byte("saltsaltsalt")),
			iterations: 10000,
			keyLength:  32,
			want:       "63086a6bd6ba533c492ed8d8453f99d723a9fbc2c427c2b848784308fc88cad3",
		},
		{
			name:       "sha256-10000-64",
			password:   []byte("password"),
			saltHex:    hex.EncodeToString([]byte("saltsaltsalt")),
			iterations: 10000,
			keyLength:  64,
			want:       "63086a6bd6ba533c492ed8d8453f99d723a9fbc2c427c2b848784308fc88cad33b33c8dad2c6f11785d281ae729f83877daa2b8afe1923791c46a8e226f92f4a",
		},
		{
			// 含 0x00 的二进制盐：验证盐是原始字节而不是十六进制字符串。
			name:       "sha256-100000-32-binary-salt",
			password:   []byte("Correct Horse Battery Staple"),
			saltHex:    "000102030405060708090a0b0c0d0e0f",
			iterations: 100000,
			keyLength:  32,
			want:       "26707b0d34414adebe877b7345529a17095c7b56bdea133de943583d1e5e2efd",
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got, err := pbkdf2x.DeriveKey(tc.password, mustHexDecode(t, tc.saltHex), tc.iterations, tc.keyLength)
			if err != nil {
				t.Fatalf("DeriveKey 返回错误: %v", err)
			}
			if hex.EncodeToString(got) != tc.want {
				t.Errorf("派生结果与已知向量不一致\n got: %s\nwant: %s", hex.EncodeToString(got), tc.want)
			}
		})
	}
}

// TestDeriveKeyWithHashKnownVector 验证 PRF 可替换：期望值同样由 Python 的
// hashlib.pbkdf2_hmac('sha512', ...) 独立生成。
func TestDeriveKeyWithHashKnownVector(t *testing.T) {
	const want = "c3d64081b5d8986fe81498e5cf0f5ed0c3bf9f011ac699ead49a7b192fe69ef239f970ba2de9827289e036f9d80fc18cc870c518d16f512cb623017df08820f8"

	got, err := pbkdf2x.DeriveKeyWithHash([]byte("password"), []byte("saltsaltsalt"), 10000, 64, sha512.New)
	if err != nil {
		t.Fatalf("DeriveKeyWithHash 返回错误: %v", err)
	}
	if hex.EncodeToString(got) != want {
		t.Errorf("PBKDF2-HMAC-SHA512 结果与已知向量不一致\n got: %s\nwant: %s", hex.EncodeToString(got), want)
	}

	// 换哈希必须得到不同结果，证明 h 真正生效而不是被写死为 SHA-256。
	sha256Key, err := pbkdf2x.DeriveKey([]byte("password"), []byte("saltsaltsalt"), 10000, 64)
	if err != nil {
		t.Fatalf("DeriveKey 返回错误: %v", err)
	}
	if bytes.Equal(sha256Key, got) {
		t.Error("SHA-512 与 SHA-256 派生结果相同，说明哈希参数未生效")
	}

	// SHA-1 仅用于兼容旧系统，这里只验证可用性与长度。
	sha1Key, err := pbkdf2x.DeriveKeyWithHash([]byte("password"), []byte("saltsaltsalt"), 10000, 20, sha1.New)
	if err != nil {
		t.Fatalf("DeriveKeyWithHash(sha1) 返回错误: %v", err)
	}
	if len(sha1Key) != 20 {
		t.Errorf("SHA-1 派生长度 = %d, 期望 20", len(sha1Key))
	}
}

// TestDeriveKeyReproducible 覆盖最核心的语义：同口令同盐同参数必须可复现。
func TestDeriveKeyReproducible(t *testing.T) {
	password := []byte("正确的马电池订书钉")
	salt := []byte("0123456789abcdef")

	k1, err := pbkdf2x.DeriveKey(password, salt, pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength)
	if err != nil {
		t.Fatalf("第一次派生失败: %v", err)
	}
	k2, err := pbkdf2x.DeriveKey(password, salt, pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength)
	if err != nil {
		t.Fatalf("第二次派生失败: %v", err)
	}
	if !bytes.Equal(k1, k2) {
		t.Fatal("相同输入两次派生结果不一致，无法复现")
	}
	if len(k1) != pbkdf2x.DefaultKeyLength {
		t.Fatalf("派生长度 = %d, 期望 %d", len(k1), pbkdf2x.DefaultKeyLength)
	}

	otherSalt, err := pbkdf2x.DeriveKey(password, []byte("fedcba9876543210"), pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength)
	if err != nil {
		t.Fatalf("换盐派生失败: %v", err)
	}
	if bytes.Equal(k1, otherSalt) {
		t.Error("不同盐得到了相同密钥")
	}

	otherPassword, err := pbkdf2x.DeriveKey([]byte("错误口令"), salt, pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength)
	if err != nil {
		t.Fatalf("换口令派生失败: %v", err)
	}
	if bytes.Equal(k1, otherPassword) {
		t.Error("不同口令得到了相同密钥")
	}

	otherIterations, err := pbkdf2x.DeriveKey(password, salt, pbkdf2x.MinIterations*2, pbkdf2x.DefaultKeyLength)
	if err != nil {
		t.Fatalf("换迭代次数派生失败: %v", err)
	}
	if bytes.Equal(k1, otherIterations) {
		t.Error("不同迭代次数得到了相同密钥")
	}
}

// TestDeriveKeyWithRandomSalt 针对被修复的核心缺陷：随机盐必须一并返回，
// 且该盐必须能复现同一个密钥。
func TestDeriveKeyWithRandomSalt(t *testing.T) {
	password := []byte("hunter2")

	key, salt, err := pbkdf2x.DeriveKeyWithRandomSalt(password, pbkdf2x.DefaultSaltLength, pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength)
	if err != nil {
		t.Fatalf("DeriveKeyWithRandomSalt 返回错误: %v", err)
	}
	if len(salt) != pbkdf2x.DefaultSaltLength {
		t.Fatalf("返回的盐长度 = %d, 期望 %d（盐必须是原始字节，不是十六进制字符串）", len(salt), pbkdf2x.DefaultSaltLength)
	}
	if len(key) != pbkdf2x.DefaultKeyLength {
		t.Fatalf("返回的密钥长度 = %d, 期望 %d", len(key), pbkdf2x.DefaultKeyLength)
	}

	// 用返回的盐复现：这是旧实现丢失盐值后无法做到的事情。
	reproduced, err := pbkdf2x.DeriveKey(password, salt, pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength)
	if err != nil {
		t.Fatalf("用返回的盐复现失败: %v", err)
	}
	if !bytes.Equal(key, reproduced) {
		t.Fatal("用 DeriveKeyWithRandomSalt 返回的盐无法复现同一密钥")
	}

	key2, salt2, err := pbkdf2x.DeriveKeyWithRandomSalt(password, pbkdf2x.DefaultSaltLength, pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength)
	if err != nil {
		t.Fatalf("第二次 DeriveKeyWithRandomSalt 失败: %v", err)
	}
	if bytes.Equal(salt, salt2) {
		t.Error("两次调用返回了相同的盐，随机盐未生效")
	}
	if bytes.Equal(key, key2) {
		t.Error("两次调用返回了相同的密钥")
	}

	// 非法 saltLength 必须报错而不是 panic。
	for _, bad := range []int{0, -1, pbkdf2x.MinSaltLength - 1, 1025, math.MaxInt} {
		if _, _, err := pbkdf2x.DeriveKeyWithRandomSalt(password, bad, pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength); err == nil {
			t.Errorf("saltLength=%d 时未返回错误", bad)
		}
	}
}

func TestVerify(t *testing.T) {
	password := []byte("s3cret-口令")
	salt := []byte("0123456789abcdef")
	key, err := pbkdf2x.DeriveKey(password, salt, pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength)
	if err != nil {
		t.Fatalf("派生失败: %v", err)
	}

	ok, err := pbkdf2x.Verify(password, salt, key, pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength)
	if err != nil {
		t.Fatalf("Verify 返回错误: %v", err)
	}
	if !ok {
		t.Error("正确口令 + 正确盐未通过校验")
	}

	cases := []struct {
		name     string
		password []byte
		salt     []byte
		expected []byte
	}{
		{"错误口令", []byte("s3cret-口令x"), salt, key},
		{"错误盐", password, []byte("fedcba9876543210"), key},
		{"密钥被篡改", password, salt, append(append([]byte{}, key[:len(key)-1]...), key[len(key)-1]^0x01)},
		{"密钥被截断", password, salt, key[:len(key)-1]},
		{"密钥被加长", password, salt, append(append([]byte{}, key...), 0x00)},
		{"空密钥", password, salt, nil},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			ok, err := pbkdf2x.Verify(tc.password, tc.salt, tc.expected, pbkdf2x.MinIterations, pbkdf2x.DefaultKeyLength)
			if err != nil {
				t.Fatalf("参数合法时 Verify 不应返回错误: %v", err)
			}
			if ok {
				t.Error("不匹配的密钥被判为匹配")
			}
		})
	}
}

// TestParameterValidation 覆盖所有非法参数路径：必须返回 error，且绝不 panic。
func TestParameterValidation(t *testing.T) {
	validSalt := bytes.Repeat([]byte{0xA5}, pbkdf2x.MinSaltLength)
	password := []byte("pw")

	cases := []struct {
		name        string
		salt        []byte
		iterations  int
		keyLength   int
		saltInvalid bool
	}{
		{"盐为 nil", nil, pbkdf2x.MinIterations, 32, true},
		{"盐为空切片", []byte{}, pbkdf2x.MinIterations, 32, true},
		{"盐过短", []byte("1234567"), pbkdf2x.MinIterations, 32, true},
		{"盐过长", bytes.Repeat([]byte{0x01}, 1025), pbkdf2x.MinIterations, 32, true},
		{"迭代次数为 0", validSalt, 0, 32, false},
		{"迭代次数为负", validSalt, -1, 32, false},
		{"迭代次数低于下限", validSalt, pbkdf2x.MinIterations - 1, 32, false},
		{"迭代次数超上限", validSalt, pbkdf2x.MaxIterations + 1, 32, false},
		{"迭代次数为 MaxInt", validSalt, math.MaxInt, 32, false},
		{"密钥长度为 0", validSalt, pbkdf2x.MinIterations, 0, false},
		{"密钥长度为负", validSalt, pbkdf2x.MinIterations, -32, false},
		{"密钥长度超上限", validSalt, pbkdf2x.MinIterations, pbkdf2x.MaxKeyLength + 1, false},
		{"密钥长度为 MaxInt", validSalt, pbkdf2x.MinIterations, math.MaxInt, false},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			key, err := pbkdf2x.DeriveKey(password, tc.salt, tc.iterations, tc.keyLength)
			if err == nil {
				t.Fatalf("DeriveKey 未返回错误，得到 %d 字节密钥", len(key))
			}
			if key != nil {
				t.Errorf("出错时密钥应为 nil，实际长度 %d", len(key))
			}

			if _, err := pbkdf2x.Verify(password, tc.salt, validSalt, tc.iterations, tc.keyLength); err == nil {
				t.Error("Verify 对非法参数未返回错误")
			}
			if _, err := pbkdf2x.DeriveKeyWithHash(password, tc.salt, tc.iterations, tc.keyLength, nil); err == nil {
				t.Error("DeriveKeyWithHash 对非法参数未返回错误")
			}
			// DeriveKeyWithRandomSalt 自行生成盐，因此仅在迭代次数/长度非法时才应报错。
			if !tc.saltInvalid {
				if _, _, err := pbkdf2x.DeriveKeyWithRandomSalt(password, pbkdf2x.DefaultSaltLength, tc.iterations, tc.keyLength); err == nil {
					t.Error("DeriveKeyWithRandomSalt 对非法参数未返回错误")
				}
			}
		})
	}

	// 合法的盐/迭代次数/长度组合但 h 为 nil：同样必须是 error 而不是 panic。
	if _, err := pbkdf2x.DeriveKeyWithHash(password, validSalt, pbkdf2x.MinIterations, 32, nil); err == nil {
		t.Error("哈希构造函数为 nil 时未返回错误")
	}
}

// TestBoundaryValues 验证边界值落在"允许"一侧时确实可用。
func TestBoundaryValues(t *testing.T) {
	salt := bytes.Repeat([]byte{0x5A}, pbkdf2x.MinSaltLength)

	// 最小迭代次数 + 最大密钥长度。
	key, err := pbkdf2x.DeriveKey([]byte("pw"), salt, pbkdf2x.MinIterations, pbkdf2x.MaxKeyLength)
	if err != nil {
		t.Fatalf("边界参数派生失败: %v", err)
	}
	if len(key) != pbkdf2x.MaxKeyLength {
		t.Errorf("密钥长度 = %d, 期望 %d", len(key), pbkdf2x.MaxKeyLength)
	}

	// 最小密钥长度。
	if _, err := pbkdf2x.DeriveKey([]byte("pw"), salt, pbkdf2x.MinIterations, 1); err != nil {
		t.Errorf("keyLength=1 应被接受: %v", err)
	}
}

func TestGenerateSalt(t *testing.T) {
	for _, length := range []int{pbkdf2x.MinSaltLength, pbkdf2x.DefaultSaltLength, 32, 64} {
		salt, err := pbkdf2x.GenerateSalt(length)
		if err != nil {
			t.Fatalf("GenerateSalt(%d) 返回错误: %v", length, err)
		}
		if len(salt) != length {
			t.Errorf("盐长度 = %d, 期望 %d", len(salt), length)
		}
	}

	a, err := pbkdf2x.GenerateSalt(pbkdf2x.DefaultSaltLength)
	if err != nil {
		t.Fatalf("GenerateSalt 返回错误: %v", err)
	}
	b, err := pbkdf2x.GenerateSalt(pbkdf2x.DefaultSaltLength)
	if err != nil {
		t.Fatalf("GenerateSalt 返回错误: %v", err)
	}
	if bytes.Equal(a, b) {
		t.Error("两次生成的盐相同，随机性有问题")
	}

	for _, length := range []int{0, -1, pbkdf2x.MinSaltLength - 1, 1025, math.MaxInt} {
		if _, err := pbkdf2x.GenerateSalt(length); err == nil {
			t.Errorf("GenerateSalt(%d) 未返回错误", length)
		}
	}
}