package kmac

import (
	"encoding/binary"

	"golang.org/x/crypto/sha3"
)

/*
算法名称: KMAC (Keccak Message Authentication Code)
算法类型: 密钥派生函数 / 消息认证码
主要特点: 基于 SHA3/Keccak 的可扩展输出函数，支持自定义字符串
应用场景: 消息认证、密钥派生、伪随机数生成
标准: NIST SP 800-185
*/

// KMAC128 使用 KMAC128 计算 MAC 或派生密钥
//
// 参数:
//   - key: 密钥
//   - data: 输入数据
//   - customization: 自定义字符串 (用于域分离，可为空)
//   - outputLen: 输出长度 (字节)
//
// 返回:
//   - MAC/派生密钥
func KMAC128(key, data []byte, customization string, outputLen int) []byte {
	// KMAC 使用 cSHAKE 作为底层，函数名固定为 "KMAC"
	h := sha3.NewCShake128([]byte("KMAC"), []byte(customization))

	// 写入编码的密钥长度和密钥
	encodedKey := encodeString(key)
	h.Write(encodedKey)

	// 写入数据
	h.Write(data)

	// 写入右编码的输出长度 (bits)
	h.Write(rightEncode(uint64(outputLen * 8)))

	// 读取输出
	output := make([]byte, outputLen)
	h.Read(output)
	return output
}

// KMAC256 使用 KMAC256 计算 MAC 或派生密钥
//
// 参数同 KMAC128，但使用更高安全级别
func KMAC256(key, data []byte, customization string, outputLen int) []byte {
	h := sha3.NewCShake256([]byte("KMAC"), []byte(customization))

	encodedKey := encodeString(key)
	h.Write(encodedKey)
	h.Write(data)
	h.Write(rightEncode(uint64(outputLen * 8)))

	output := make([]byte, outputLen)
	h.Read(output)
	return output
}

// DeriveKey 使用 KMAC256 派生密钥
//
// 参数:
//   - masterKey: 主密钥
//   - context: 上下文/标签 (用于域分离)
//   - keyLen: 输出密钥长度
//
// 返回:
//   - 派生的密钥
func DeriveKey(masterKey, context []byte, keyLen int) []byte {
	return KMAC256(masterKey, context, "KeyDerivation", keyLen)
}

// DeriveKeyWithLabel 使用标签派生密钥
//
// 参数:
//   - masterKey: 主密钥
//   - label: 字符串标签
//   - context: 二进制上下文
//   - keyLen: 输出长度
func DeriveKeyWithLabel(masterKey []byte, label string, context []byte, keyLen int) []byte {
	// 组合标签和上下文
	data := append([]byte(label), context...)
	return KMAC256(masterKey, data, "LabeledKDF", keyLen)
}

// MAC128 计算 128-bit MAC
func MAC128(key, message []byte) []byte {
	return KMAC128(key, message, "", 16)
}

// MAC256 计算 256-bit MAC
func MAC256(key, message []byte) []byte {
	return KMAC256(key, message, "", 32)
}

// Verify 验证 MAC
//
// 使用常量时间比较防止时序攻击
func Verify(key, message, expectedMAC []byte) bool {
	if len(expectedMAC) == 0 {
		return false
	}
	computedMAC := KMAC256(key, message, "", len(expectedMAC))
	return constantTimeCompare(computedMAC, expectedMAC)
}

// constantTimeCompare 常量时间比较
func constantTimeCompare(a, b []byte) bool {
	if len(a) != len(b) {
		return false
	}
	var result byte
	for i := range a {
		result |= a[i] ^ b[i]
	}
	return result == 0
}

// KMACXOF128 可扩展输出函数模式 (无限输出)
// 返回一个 io.Reader，可按需读取任意长度的输出
func KMACXOF128(key, data []byte, customization string) *sha3.ShakeHash {
	h := sha3.NewCShake128([]byte("KMAC"), []byte(customization))

	encodedKey := encodeString(key)
	h.Write(encodedKey)
	h.Write(data)
	// XOF 模式：输出长度编码为 0
	h.Write(rightEncode(0))

	return &h
}

// KMACXOF256 可扩展输出函数模式 (无限输出)
func KMACXOF256(key, data []byte, customization string) *sha3.ShakeHash {
	h := sha3.NewCShake256([]byte("KMAC"), []byte(customization))

	encodedKey := encodeString(key)
	h.Write(encodedKey)
	h.Write(data)
	h.Write(rightEncode(0))

	return &h
}

// encodeString 按 SP 800-185 编码字符串
// encode_string(S) = left_encode(len(S)) || S
func encodeString(s []byte) []byte {
	encoded := leftEncode(uint64(len(s) * 8))
	return append(encoded, s...)
}

// leftEncode 左编码整数
// 用于编码字符串长度
func leftEncode(x uint64) []byte {
	if x == 0 {
		return []byte{1, 0}
	}

	var buf [9]byte
	n := 0
	for temp := x; temp > 0; temp >>= 8 {
		n++
	}
	for i := n; i > 0; i-- {
		buf[i] = byte(x)
		x >>= 8
	}
	buf[0] = byte(n)
	return buf[:n+1]
}

// rightEncode 右编码整数
// 用于编码输出长度
func rightEncode(x uint64) []byte {
	if x == 0 {
		return []byte{0, 1}
	}

	var buf [9]byte
	n := 0
	for temp := x; temp > 0; temp >>= 8 {
		n++
		buf[8-n+1] = byte(temp)
	}
	buf[9-n] = buf[9-n] // 保持正确位置
	result := buf[9-n:]
	return append(result, byte(n))
}

// DeriveMultipleKeys 派生多个密钥
func DeriveMultipleKeys(masterKey []byte, labels []string, keyLen int) [][]byte {
	keys := make([][]byte, len(labels))
	for i, label := range labels {
		keys[i] = DeriveKeyWithLabel(masterKey, label, nil, keyLen)
	}
	return keys
}

// PRF 伪随机函数
// 可用于生成确定性随机数
func PRF(key []byte, counter uint64, outputLen int) []byte {
	counterBytes := make([]byte, 8)
	binary.BigEndian.PutUint64(counterBytes, counter)
	return KMAC256(key, counterBytes, "PRF", outputLen)
}