package hkdfx_test

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

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

// seq 生成从 start 开始、长度为 n 的连续字节，用于构造 RFC 5869 的测试输入。
func seq(start byte, n int) []byte {
	b := make([]byte, n)
	for i := range b {
		b[i] = start + byte(i)
	}
	return b
}

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

// TestDeriveKeyRFC5869Vectors 使用 RFC 5869 附录 A 的官方向量做已知答案测试。
// 这些向量同时覆盖了 SHA-256 与 SHA-1、空盐以及 82 字节的长输出。
func TestDeriveKeyRFC5869Vectors(t *testing.T) {
	cases := []struct {
		name      string
		secret    []byte
		salt      []byte
		info      []byte
		keyLength int
		useSHA1   bool
		want      string
	}{
		{
			name:      "RFC5869-A.1 SHA-256 基本向量",
			secret:    bytes.Repeat([]byte{0x0b}, 22),
			salt:      seq(0x00, 13),
			info:      seq(0xf0, 10),
			keyLength: 42,
			want:      "3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865",
		},
		{
			name:      "RFC5869-A.2 SHA-256 长输入长输出",
			secret:    seq(0x00, 80),
			salt:      seq(0x60, 80),
			info:      seq(0xb0, 80),
			keyLength: 82,
			want:      "b11e398dc80327a1c8e7f78c596a49344f012eda2d4efad8a050cc4c19afa97c59045a99cac7827271cb41c65e590e09da3275600c2f09b8367793a9aca3db71cc30c58179ec3e87c14c01d5c1f3434f1d87",
		},
		{
			name:      "RFC5869-A.3 SHA-256 空盐空 info",
			secret:    bytes.Repeat([]byte{0x0b}, 22),
			salt:      nil,
			info:      nil,
			keyLength: 42,
			want:      "8da4e775a563c18f715f802a063c5a31b8a11f5c5ee1879ec3454e5f3c738d2d9d201395faa4b61a96c8",
		},
		{
			name:      "RFC5869-A.4 SHA-1 基本向量",
			secret:    bytes.Repeat([]byte{0x0b}, 11),
			salt:      seq(0x00, 13),
			info:      seq(0xf0, 10),
			keyLength: 42,
			useSHA1:   true,
			want:      "085a01ea1b10f36933068b56efa5ad81a4f14b822f5b091568a9cdd4f155fda2c22e422478d305f3f896",
		},
		{
			name:      "RFC5869-A.5 SHA-1 长输入长输出",
			secret:    seq(0x00, 80),
			salt:      seq(0x60, 80),
			info:      seq(0xb0, 80),
			keyLength: 82,
			useSHA1:   true,
			want:      "0bd770a74d1160f7c9f12cd5912a06ebff6adcae899d92191fe4305673ba2ffe8fa3f1a4e5ad79f3f334b3b202b2173c486ea37ce3d397ed034c7f9dfeb15c5e927336d0441f4c4300e2cff0d0900b52d3b4",
		},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			var (
				got []byte
				err error
			)
			if tc.useSHA1 {
				got, err = hkdfx.DeriveKeyWithHash(tc.secret, tc.salt, tc.info, tc.keyLength, sha1.New)
			} else {
				got, err = hkdfx.DeriveKey(tc.secret, tc.salt, tc.info, tc.keyLength)
			}
			if err != nil {
				t.Fatalf("派生失败: %v", err)
			}
			if hex.EncodeToString(got) != tc.want {
				t.Errorf("与 RFC 5869 向量不一致\n got: %s\nwant: %s", hex.EncodeToString(got), tc.want)
			}
		})
	}
}

func TestDeriveKeyProperties(t *testing.T) {
	secret := mustHexDecode(t, "000102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f")
	salt := []byte("unique-salt")
	info := []byte("aes-key")

	k1, err := hkdfx.DeriveKey(secret, salt, info, hkdfx.DefaultKeyLength)
	if err != nil {
		t.Fatalf("派生失败: %v", err)
	}
	k2, err := hkdfx.DeriveKey(secret, salt, info, hkdfx.DefaultKeyLength)
	if err != nil {
		t.Fatalf("派生失败: %v", err)
	}
	if !bytes.Equal(k1, k2) {
		t.Error("相同输入两次派生结果不一致")
	}
	if len(k1) != hkdfx.DefaultKeyLength {
		t.Errorf("派生长度 = %d, 期望 %d", len(k1), hkdfx.DefaultKeyLength)
	}

	// 默认入口必须等价于显式传 SHA-256。
	explicit, err := hkdfx.DeriveKeyWithHash(secret, salt, info, hkdfx.DefaultKeyLength, sha256.New)
	if err != nil {
		t.Fatalf("派生失败: %v", err)
	}
	if !bytes.Equal(k1, explicit) {
		t.Error("DeriveKey 与 DeriveKeyWithHash(sha256) 结果不一致")
	}

	otherInfo, err := hkdfx.DeriveKey(secret, salt, []byte("hmac-key"), hkdfx.DefaultKeyLength)
	if err != nil {
		t.Fatalf("派生失败: %v", err)
	}
	if bytes.Equal(k1, otherInfo) {
		t.Error("不同 info 得到了相同密钥（密钥分离失效）")
	}

	otherSalt, err := hkdfx.DeriveKey(secret, []byte("another-salt"), info, hkdfx.DefaultKeyLength)
	if err != nil {
		t.Fatalf("派生失败: %v", err)
	}
	if bytes.Equal(k1, otherSalt) {
		t.Error("不同 salt 得到了相同密钥")
	}

	otherSecret, err := hkdfx.DeriveKey(append([]byte{0x00}, secret...), salt, info, hkdfx.DefaultKeyLength)
	if err != nil {
		t.Fatalf("派生失败: %v", err)
	}
	if bytes.Equal(k1, otherSecret) {
		t.Error("不同 secret 得到了相同密钥")
	}

	// 前缀性质：先派生 64 字节，再取前 32 字节应与直接派生 32 字节一致。
	long, err := hkdfx.DeriveKey(secret, salt, info, 64)
	if err != nil {
		t.Fatalf("派生失败: %v", err)
	}
	if !bytes.Equal(long[:32], k1) {
		t.Error("HKDF 输出不满足前缀一致性")
	}
}

// TestDeriveKeyNilSaltEqualsZeroSalt 验证 nil 盐按 RFC 5869 等价于 HashLen 个零字节。
func TestDeriveKeyNilSaltEqualsZeroSalt(t *testing.T) {
	secret := []byte("high-entropy-secret-material")
	info := []byte("context")

	withNil, err := hkdfx.DeriveKey(secret, nil, info, 32)
	if err != nil {
		t.Fatalf("nil 盐派生失败: %v", err)
	}
	withEmpty, err := hkdfx.DeriveKey(secret, []byte{}, info, 32)
	if err != nil {
		t.Fatalf("空盐派生失败: %v", err)
	}
	withZeros, err := hkdfx.DeriveKey(secret, make([]byte, sha256.Size), info, 32)
	if err != nil {
		t.Fatalf("零盐派生失败: %v", err)
	}
	if !bytes.Equal(withNil, withEmpty) || !bytes.Equal(withNil, withZeros) {
		t.Error("nil 盐、空盐与 HashLen 个零字节盐的派生结果不一致")
	}
}

func TestParameterValidation(t *testing.T) {
	secret := []byte("high-entropy-secret-material")

	// secret 必须非空。
	for _, s := range [][]byte{nil, {}} {
		if _, err := hkdfx.DeriveKey(s, nil, nil, 32); err == nil {
			t.Error("空 secret 未返回错误")
		}
		if _, err := hkdfx.DeriveKeyWithHash(s, nil, nil, 32, sha256.New); err == nil {
			t.Error("空 secret 未返回错误")
		}
	}

	// 长度非法。
	for _, length := range []int{0, -1, hkdfx.MaxKeyLength + 1, math.MaxInt} {
		if _, err := hkdfx.DeriveKey(secret, nil, nil, length); err == nil {
			t.Errorf("keyLength=%d 未返回错误", length)
		}
	}

	// 哈希为 nil。
	if _, err := hkdfx.DeriveKeyWithHash(secret, nil, nil, 32, nil); err == nil {
		t.Error("哈希构造函数为 nil 时未返回错误")
	}
}

// TestKeyLengthLimits 验证上限按 255*HashLen 计算，并允许等于上限。
func TestKeyLengthLimits(t *testing.T) {
	secret := []byte("high-entropy-secret-material")

	cases := []struct {
		name      string
		h         func() hash.Hash
		maxLength int
	}{
		{"sha256", sha256.New, 255 * sha256.Size},
		{"sha512", sha512.New, 255 * sha512.Size},
		{"sha1", sha1.New, 255 * sha1.Size},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			key, err := hkdfx.DeriveKeyWithHash(secret, []byte("salt"), []byte("info"), tc.maxLength, tc.h)
			if err != nil {
				t.Fatalf("取上限 %d 时不应报错: %v", tc.maxLength, err)
			}
			if len(key) != tc.maxLength {
				t.Errorf("派生长度 = %d, 期望 %d", len(key), tc.maxLength)
			}
			if _, err := hkdfx.DeriveKeyWithHash(secret, []byte("salt"), []byte("info"), tc.maxLength+1, tc.h); err == nil {
				t.Errorf("超过哈希上限 %d 时未返回错误", tc.maxLength)
			}
		})
	}

	// 默认 SHA-256 入口的常量上限必须与 255*32 一致。
	if hkdfx.MaxKeyLength != 255*sha256.Size {
		t.Errorf("MaxKeyLength = %d, 期望 %d", hkdfx.MaxKeyLength, 255*sha256.Size)
	}
}