package sha3x_test

import (
	"bytes"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"hash"
	"io"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"code.yun.ink/pkg/encryptx/sha3x"
	"golang.org/x/crypto/sha3"
)

// 测试向量来源：NIST FIPS 202 / SHA-3 标准示例，并经 Python 3 hashlib 独立复算确认。
var vectors = []struct {
	name      string
	newHash   func() hash.Hash
	sum       func([]byte) []byte
	hex       func([]byte) string
	b64       func([]byte) string
	sumReader func(io.Reader) ([]byte, error)
	hexReader func(io.Reader) (string, error)
	hexFile   func(string) (string, error)
	size      int
	empty     string
	abc       string
}{
	{
		name:      "SHA3-224",
		newHash:   sha3x.New224,
		sum:       sha3x.Sum224,
		hex:       sha3x.Hex224,
		b64:       sha3x.Base64Sum224,
		sumReader: sha3x.Sum224Reader,
		hexReader: sha3x.Hex224Reader,
		hexFile:   sha3x.HexFile224,
		size:      sha3x.Size224,
		empty:     "6b4e03423667dbb73b6e15454f0eb1abd4597f9a1b078e3f5b5a6bc7",
		abc:       "e642824c3f8cf24ad09234ee7d3c766fc9a3a5168d0c94ad73b46fdf",
	},
	{
		name:      "SHA3-256",
		newHash:   sha3x.New256,
		sum:       sha3x.Sum256,
		hex:       sha3x.Hex256,
		b64:       sha3x.Base64Sum256,
		sumReader: sha3x.Sum256Reader,
		hexReader: sha3x.Hex256Reader,
		hexFile:   sha3x.HexFile256,
		size:      sha3x.Size256,
		empty:     "a7ffc6f8bf1ed76651c14756a061d662f580ff4de43b49fa82d80a4b80f8434a",
		abc:       "3a985da74fe225b2045c172d6bd390bd855f086e3e9d525b46bfe24511431532",
	},
	{
		name:      "SHA3-384",
		newHash:   sha3x.New384,
		sum:       sha3x.Sum384,
		hex:       sha3x.Hex384,
		b64:       sha3x.Base64Sum384,
		sumReader: sha3x.Sum384Reader,
		hexReader: sha3x.Hex384Reader,
		hexFile:   sha3x.HexFile384,
		size:      sha3x.Size384,
		empty:     "0c63a75b845e4f7d01107d852e4c2485c51a50aaaa94fc61995e71bbee983a2ac3713831264adb47fb6bd1e058d5f004",
		abc:       "ec01498288516fc926459f58e2c6ad8df9b473cb0fc08c2596da7cf0e49be4b298d88cea927ac7f539f1edf228376d25",
	},
	{
		name:      "SHA3-512",
		newHash:   sha3x.New512,
		sum:       sha3x.Sum512,
		hex:       sha3x.Hex512,
		b64:       sha3x.Base64Sum512,
		sumReader: sha3x.Sum512Reader,
		hexReader: sha3x.Hex512Reader,
		hexFile:   sha3x.HexFile512,
		size:      sha3x.Size512,
		empty:     "a69f73cca23a9ac5c8b567dc185a756e97c982164fe25859e0d1dcc1475c80a615b2123af1f5f94c11e3e9402c3ac558f500199d95b6d3e301758586281dcd26",
		abc:       "b751850b1a57168a5693cd924b6b096e08f621827444f70d884f5d0240d2712e10e116e9192af3c91a7ec57647e3934057340b4cf408d5a56592f8274eec53f0",
	},
}

type failingReader struct{}

func (failingReader) Read([]byte) (int, error) { return 0, errors.New("模拟读取失败") }

func TestConstants(t *testing.T) {
	sizes := map[string][2]int{
		"SHA3-224": {sha3x.Size224, sha3x.Rate224},
		"SHA3-256": {sha3x.Size256, sha3x.Rate256},
		"SHA3-384": {sha3x.Size384, sha3x.Rate384},
		"SHA3-512": {sha3x.Size512, sha3x.Rate512},
	}
	want := map[string][2]int{
		"SHA3-224": {28, 144},
		"SHA3-256": {32, 136},
		"SHA3-384": {48, 104},
		"SHA3-512": {64, 72},
	}
	for name, got := range sizes {
		if got != want[name] {
			t.Errorf("%s: Size/Rate = %v, 期望 %v", name, got, want[name])
		}
	}
	if sha3x.MaxShakeOutputLength != 1<<20 {
		t.Errorf("MaxShakeOutputLength = %d, 期望 %d", sha3x.MaxShakeOutputLength, 1<<20)
	}
}

func TestVectors(t *testing.T) {
	for _, v := range vectors {
		t.Run(v.name, func(t *testing.T) {
			cases := []struct {
				name  string
				input string
				want  string
			}{
				{"空串", "", v.empty},
				{"abc", "abc", v.abc},
			}
			for _, tc := range cases {
				t.Run(tc.name, func(t *testing.T) {
					data := []byte(tc.input)
					sum := v.sum(data)
					if len(sum) != v.size {
						t.Fatalf("Sum 长度 = %d, 期望 %d", len(sum), v.size)
					}
					if got := hex.EncodeToString(sum); got != tc.want {
						t.Errorf("Sum = %s, 期望 %s", got, tc.want)
					}
					if got := v.hex(data); got != tc.want {
						t.Errorf("Hex = %s, 期望 %s", got, tc.want)
					}
					if want := base64.StdEncoding.EncodeToString(sum); v.b64(data) != want {
						t.Errorf("Base64 = %s, 期望 %s", v.b64(data), want)
					}
					h := v.newHash()
					if _, err := h.Write(data); err != nil {
						t.Fatalf("Write 失败: %v", err)
					}
					if got := hex.EncodeToString(h.Sum(nil)); got != tc.want {
						t.Errorf("New+Write = %s, 期望 %s", got, tc.want)
					}
				})
			}
		})
	}
}

func TestDefaultAliasesMatch256(t *testing.T) {
	data := []byte("encryptx alias check")
	if !bytes.Equal(sha3x.Sum(data), sha3x.Sum256(data)) {
		t.Error("Sum 应等价于 Sum256")
	}
	if sha3x.Hex(data) != sha3x.Hex256(data) {
		t.Error("Hex 应等价于 Hex256")
	}
	if sha3x.Base64(data) != sha3x.Base64Sum256(data) {
		t.Error("Base64 应等价于 Base64Sum256")
	}
	if h := sha3x.New(); h.Size() != sha3x.Size256 {
		t.Errorf("New().Size() = %d, 期望 %d", h.Size(), sha3x.Size256)
	}
	if got, err := sha3x.SumReader(strings.NewReader(string(data))); err != nil {
		t.Fatalf("SumReader 失败: %v", err)
	} else if !bytes.Equal(got, sha3x.Sum256(data)) {
		t.Error("SumReader 应等价于 Sum256Reader")
	}
	if got, err := sha3x.HexReader(strings.NewReader(string(data))); err != nil {
		t.Fatalf("HexReader 失败: %v", err)
	} else if got != sha3x.Hex256(data) {
		t.Error("HexReader 应等价于 Hex256Reader")
	}
}

func TestStreamingMatchesSum(t *testing.T) {
	data := []byte(strings.Repeat("encryptx-sha3-", 200))
	for _, v := range vectors {
		t.Run(v.name, func(t *testing.T) {
			h := v.newHash()
			for i := 0; i < len(data); i += 17 {
				end := i + 17
				if end > len(data) {
					end = len(data)
				}
				if _, err := h.Write(data[i:end]); err != nil {
					t.Fatalf("Write 失败: %v", err)
				}
			}
			if got, want := h.Sum(nil), v.sum(data); !bytes.Equal(got, want) {
				t.Errorf("分块写入 = %x, 期望 %x", got, want)
			}
		})
	}
}

func TestReaderConsistency(t *testing.T) {
	const s = "encryptx 一致性检查"
	for _, v := range vectors {
		t.Run(v.name, func(t *testing.T) {
			sum, err := v.sumReader(strings.NewReader(s))
			if err != nil {
				t.Fatalf("SumReader 失败: %v", err)
			}
			if want := v.sum([]byte(s)); !bytes.Equal(sum, want) {
				t.Errorf("SumReader = %x, 期望 %x", sum, want)
			}
			hexStr, err := v.hexReader(strings.NewReader(s))
			if err != nil {
				t.Fatalf("HexReader 失败: %v", err)
			}
			if want := v.hex([]byte(s)); hexStr != want {
				t.Errorf("HexReader = %s, 期望 %s", hexStr, want)
			}
		})
	}
}

func TestReaderErrors(t *testing.T) {
	readers := []func(io.Reader) ([]byte, error){
		sha3x.Sum224Reader, sha3x.Sum256Reader, sha3x.Sum384Reader, sha3x.Sum512Reader, sha3x.SumReader,
	}
	for i, fn := range readers {
		if _, err := fn(nil); err == nil {
			t.Errorf("第 %d 个读取函数：nil 读取器应返回 error", i)
		}
		if _, err := fn(failingReader{}); err == nil {
			t.Errorf("第 %d 个读取函数：读取失败时应返回 error", i)
		}
	}
	hexReaders := []func(io.Reader) (string, error){
		sha3x.Hex224Reader, sha3x.Hex256Reader, sha3x.Hex384Reader, sha3x.Hex512Reader, sha3x.HexReader,
	}
	for i, fn := range hexReaders {
		if _, err := fn(nil); err == nil {
			t.Errorf("第 %d 个十六进制读取函数：nil 读取器应返回 error", i)
		}
		if _, err := fn(failingReader{}); err == nil {
			t.Errorf("第 %d 个十六进制读取函数：读取失败时应返回 error", i)
		}
	}
}

func TestHexFile(t *testing.T) {
	dir := t.TempDir()
	content := []byte("文件内容 file content\n")
	path := filepath.Join(dir, "sample.bin")
	if err := os.WriteFile(path, content, 0o600); err != nil {
		t.Fatalf("写入测试文件失败: %v", err)
	}
	for _, v := range vectors {
		t.Run(v.name, func(t *testing.T) {
			got, err := v.hexFile(path)
			if err != nil {
				t.Fatalf("HexFile 失败: %v", err)
			}
			if want := v.hex(content); got != want {
				t.Errorf("HexFile = %s, 期望 %s", got, want)
			}
		})
	}
	if got, err := sha3x.HexFile(path); err != nil {
		t.Fatalf("HexFile 失败: %v", err)
	} else if want := sha3x.Hex256(content); got != want {
		t.Errorf("HexFile = %s, 期望 %s", got, want)
	}
}

func TestHexFileErrors(t *testing.T) {
	missing := filepath.Join(t.TempDir(), "不存在.bin")
	files := []func(string) (string, error){
		sha3x.HexFile224, sha3x.HexFile256, sha3x.HexFile384, sha3x.HexFile512, sha3x.HexFile,
	}
	for i, fn := range files {
		if _, err := fn(missing); err == nil {
			t.Errorf("第 %d 个文件函数：不存在的文件应返回 error", i)
		}
		if _, err := fn(""); err == nil {
			t.Errorf("第 %d 个文件函数：空路径应返回 error", i)
		}
		if _, err := fn(t.TempDir()); err == nil {
			t.Errorf("第 %d 个文件函数：目录路径应返回 error", i)
		}
	}
}

// 测试向量来源：NIST FIPS 202 SHAKE 示例，并经 Python 3 hashlib 独立复算确认。
func TestShakeVectors(t *testing.T) {
	cases := []struct {
		name    string
		fn      func([]byte, int) ([]byte, error)
		fnHex   func([]byte, int) (string, error)
		input   string
		n       int
		wantHex string
	}{
		{"SHAKE128 空串 16 字节", sha3x.Shake128, sha3x.Shake128Hex, "", 16, "7f9c2ba4e88f827d616045507605853e"},
		{"SHAKE128 abc 16 字节", sha3x.Shake128, sha3x.Shake128Hex, "abc", 16, "5881092dd818bf5cf8a3ddb793fbcba7"},
		{"SHAKE256 空串 32 字节", sha3x.Shake256, sha3x.Shake256Hex, "", 32, "46b9dd2b0ba88d13233b3feb743eeb243fcd52ea62b81b82b50c27646ed5762f"},
		{"SHAKE256 abc 32 字节", sha3x.Shake256, sha3x.Shake256Hex, "abc", 32, "483366601360a8771c6863080cc4114d8db44530f8f1e1ee4f94ea37e78b5739"},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			out, err := tc.fn([]byte(tc.input), tc.n)
			if err != nil {
				t.Fatalf("计算失败: %v", err)
			}
			if len(out) != tc.n {
				t.Fatalf("输出长度 = %d, 期望 %d", len(out), tc.n)
			}
			if got := hex.EncodeToString(out); got != tc.wantHex {
				t.Errorf("输出 = %s, 期望 %s", got, tc.wantHex)
			}
			gotHex, err := tc.fnHex([]byte(tc.input), tc.n)
			if err != nil {
				t.Fatalf("计算十六进制失败: %v", err)
			}
			if gotHex != tc.wantHex {
				t.Errorf("十六进制输出 = %s, 期望 %s", gotHex, tc.wantHex)
			}
		})
	}
}

func TestShakeLengthValidation(t *testing.T) {
	bad := []int{0, -1, -1024, sha3x.MaxShakeOutputLength + 1, 1 << 30}
	for _, n := range bad {
		if _, err := sha3x.Shake128([]byte("x"), n); err == nil {
			t.Errorf("Shake128 长度 %d 应返回 error", n)
		}
		if _, err := sha3x.Shake256([]byte("x"), n); err == nil {
			t.Errorf("Shake256 长度 %d 应返回 error", n)
		}
		if _, err := sha3x.Shake128Hex([]byte("x"), n); err == nil {
			t.Errorf("Shake128Hex 长度 %d 应返回 error", n)
		}
		if _, err := sha3x.Shake256Hex([]byte("x"), n); err == nil {
			t.Errorf("Shake256Hex 长度 %d 应返回 error", n)
		}
		if _, err := sha3x.Shake128Reader(strings.NewReader("x"), n); err == nil {
			t.Errorf("Shake128Reader 长度 %d 应返回 error", n)
		}
		if _, err := sha3x.Shake256Reader(strings.NewReader("x"), n); err == nil {
			t.Errorf("Shake256Reader 长度 %d 应返回 error", n)
		}
	}
	if _, err := sha3x.Shake128Reader(nil, 16); err == nil {
		t.Error("nil 读取器应返回 error")
	}
	if _, err := sha3x.Shake256Reader(nil, 16); err == nil {
		t.Error("nil 读取器应返回 error")
	}
	if _, err := sha3x.Shake128Reader(failingReader{}, 16); err == nil {
		t.Error("读取失败时应返回 error")
	}
}

func TestShakeOutputProperties(t *testing.T) {
	full, err := sha3x.Shake128([]byte("xyz"), 64)
	if err != nil {
		t.Fatalf("计算失败: %v", err)
	}
	short, err := sha3x.Shake128([]byte("xyz"), 32)
	if err != nil {
		t.Fatalf("计算失败: %v", err)
	}
	if !bytes.Equal(short, full[:32]) {
		t.Error("XOF 前缀性质不成立：短输出应是长输出的前缀")
	}
	other, err := sha3x.Shake128([]byte("xy"), 64)
	if err != nil {
		t.Fatalf("计算失败: %v", err)
	}
	if bytes.Equal(full, other) {
		t.Error("不同输入的 SHAKE 输出不应相同")
	}

	max, err := sha3x.Shake256([]byte("x"), sha3x.MaxShakeOutputLength)
	if err != nil {
		t.Fatalf("上限长度计算失败: %v", err)
	}
	if len(max) != sha3x.MaxShakeOutputLength {
		t.Errorf("上限长度输出 = %d, 期望 %d", len(max), sha3x.MaxShakeOutputLength)
	}
}

func TestShakeReaderConsistency(t *testing.T) {
	const s = "encryptx shake reader"
	data := []byte(s)
	got, err := sha3x.Shake128Reader(strings.NewReader(s), 48)
	if err != nil {
		t.Fatalf("Shake128Reader 失败: %v", err)
	}
	if want, err := sha3x.Shake128(data, 48); err != nil {
		t.Fatalf("Shake128 失败: %v", err)
	} else if !bytes.Equal(got, want) {
		t.Errorf("Shake128Reader = %x, 期望 %x", got, want)
	}

	got256, err := sha3x.Shake256Reader(strings.NewReader(s), 48)
	if err != nil {
		t.Fatalf("Shake256Reader 失败: %v", err)
	}
	if want, err := sha3x.Shake256(data, 48); err != nil {
		t.Fatalf("Shake256 失败: %v", err)
	} else if !bytes.Equal(got256, want) {
		t.Errorf("Shake256Reader = %x, 期望 %x", got256, want)
	}
}

func TestShakeStreaming(t *testing.T) {
	data := []byte(strings.Repeat("encryptx-xof-", 64))
	cases := []struct {
		name    string
		newXOF  func() sha3.ShakeHash
		wantOut func([]byte, int) ([]byte, error)
	}{
		{"SHAKE128", sha3x.NewShake128, sha3x.Shake128},
		{"SHAKE256", sha3x.NewShake256, sha3x.Shake256},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			h := tc.newXOF()
			for i := 0; i < len(data); i += 5 {
				end := i + 5
				if end > len(data) {
					end = len(data)
				}
				if _, err := h.Write(data[i:end]); err != nil {
					t.Fatalf("Write 失败: %v", err)
				}
			}

			out := make([]byte, 40)
			if _, err := io.ReadFull(h, out); err != nil {
				t.Fatalf("读取 XOF 输出失败: %v", err)
			}
			want, err := tc.wantOut(data, 40)
			if err != nil {
				t.Fatalf("计算期望值失败: %v", err)
			}
			if !bytes.Equal(out, want) {
				t.Errorf("流式 XOF = %x, 期望 %x", out, want)
			}

			// 同一 XOF 实例续读时应保持输出连续。
			rest := make([]byte, 24)
			if _, err := io.ReadFull(h, rest); err != nil {
				t.Fatalf("续读 XOF 输出失败: %v", err)
			}
			full, err := tc.wantOut(data, 64)
			if err != nil {
				t.Fatalf("计算期望值失败: %v", err)
			}
			if !bytes.Equal(rest, full[40:]) {
				t.Errorf("续读输出 = %x, 期望 %x", rest, full[40:])
			}
		})
	}
}

func TestLargeInput(t *testing.T) {
	data := largeData(1 << 20)
	for _, v := range vectors {
		t.Run(v.name, func(t *testing.T) {
			want := v.hex(data)
			if got := v.hex(data); got != want {
				t.Errorf("重复计算结果不一致: %s vs %s", got, want)
			}
			got, err := v.hexReader(bytes.NewReader(data))
			if err != nil {
				t.Fatalf("HexReader 失败: %v", err)
			}
			if got != want {
				t.Errorf("HexReader = %s, 期望 %s", got, want)
			}
			h := v.newHash()
			if _, err := h.Write(data[:1<<19]); err != nil {
				t.Fatalf("Write 失败: %v", err)
			}
			if _, err := h.Write(data[1<<19:]); err != nil {
				t.Fatalf("Write 失败: %v", err)
			}
			if got := hex.EncodeToString(h.Sum(nil)); got != want {
				t.Errorf("1MiB 分块 = %s, 期望 %s", got, want)
			}
		})
	}
}

// largeData 生成确定性的 n 字节数据。
func largeData(n int) []byte {
	b := make([]byte, n)
	for i := range b {
		b[i] = byte(i*31 + i/7)
	}
	return b
}