package pbkd

// pbkdf 实际应为pbkdf2

import (
	"crypto/hmac"
	"encoding/binary"
	"errors"
	"hash"

	"golang.org/x/crypto/pbkdf2"
	"xdx.jelly/xgcl/internal/xor"
	"xdx.jelly/xgcl/sm/sm3"
)

const (
	MinimumCount         = 1024
	MildCount            = 100000
	RecommendCount       = 10000000
	two32m1        int64 = 0xFFFFFFFF
)

var errDKLenTooLong = errors.New("dkLen too long")

// f 实现错误，正确应为f2
// 虎符中用的是f
func f(prfer prfer, salt_i []byte, count int) ([]byte, error) {
	u, err := prfer.prf(salt_i)
	if err != nil {
		return nil, err
	}
	for i := 2; i <= count; i++ {
		if out, err := prfer.prf(u); err != nil {
			return nil, err
		} else {
			xor.XorBytes(u, u, out)
		}
	}
	return u, nil
}
func kdf(prfer prfer, salt []byte, count int, dkLen int64) ([]byte, error) {
	hLen := prfer.hLen()
	if dkLen > two32m1*hLen {
		return nil, errDKLenTooLong
	}
	n := (dkLen + hLen - 1) / hLen
	dk := make([]byte, 0, n*hLen)
	sLen := len(salt)
	salt_i := make([]byte, sLen+4)
	copy(salt_i, salt)

	for i := uint32(1); int64(i) <= n; i++ {
		binary.BigEndian.PutUint32(salt_i[sLen:], i)
		if out, err := f(prfer, salt_i, count); err != nil {
			return nil, err
		} else {
			dk = append(dk, out...)
		}
	}
	return dk[:dkLen], nil
}

func f2(prfer prfer, salt_i []byte, count int) ([]byte, error) {
	u, err := prfer.prf(salt_i)
	if err != nil {
		return nil, err
	}
	T := append([]byte(nil), u...)
	for i := 2; i <= count; i++ {
		if u, err = prfer.prf(u); err != nil {
			return nil, err
		} else {
			xor.XorBytes(T, T, u)
		}
	}
	return T, nil
}

func kdf2(prfer prfer, salt []byte, count int, dkLen int64) ([]byte, error) {
	hLen := prfer.hLen()
	if dkLen > two32m1*hLen {
		return nil, errDKLenTooLong
	}
	n := (dkLen + hLen - 1) / hLen
	dk := make([]byte, 0, n*hLen)
	sLen := len(salt)
	salt_i := make([]byte, sLen+4)
	copy(salt_i, salt)

	for i := uint32(1); int64(i) <= n; i++ {
		binary.BigEndian.PutUint32(salt_i[sLen:], i)
		if out, err := f2(prfer, salt_i, count); err != nil {
			return nil, err
		} else {
			dk = append(dk, out...)
		}
	}
	return dk[:dkLen], nil
}

// PbkdfWithHmacSm3 计算口令导出密钥。遵循GM/T 0091，使用SM3_HMAC作为PRF。
func PbkdfWithHmacSm3(password, salt []byte, count int, dkLen int64) ([]byte, error) {
	return pbkdf2.Key([]byte(password), salt, count, int(dkLen), sm3.New), nil
	// return kdf2(newPrfHmacSm3(password), salt, count, dkLen)
	// return key([]byte(password), salt, count, int(dkLen), sm3.New), nil
}

// key 按GM/T 0091计算导出密钥。
func key(password, salt []byte, iter, keyLen int, h func() hash.Hash) []byte {
	prf := hmac.New(h, password)
	hashLen := prf.Size()
	numBlocks := (keyLen + hashLen - 1) / hashLen

	var buf [4]byte
	dk := make([]byte, 0, numBlocks*hashLen)
	U := make([]byte, hashLen)
	for block := 1; block <= numBlocks; block++ {
		// N.B.: || means concatenation, ^ means XOR
		// for each block T_i = U_1 ^ U_2 ^ ... ^ U_iter
		// U_1 = PRF(password, salt || uint(i))
		prf.Reset()
		prf.Write(salt)
		buf[0] = byte(block >> 24)
		buf[1] = byte(block >> 16)
		buf[2] = byte(block >> 8)
		buf[3] = byte(block)
		prf.Write(buf[:4])
		dk = prf.Sum(dk)
		T := dk[len(dk)-hashLen:]
		copy(U, T)

		// U_n = PRF(password, U_(n-1))
		for n := 2; n <= iter; n++ {
			prf.Reset()
			prf.Write(U)
			U = U[:0]
			U = prf.Sum(U)
			for x := range U {
				T[x] ^= U[x]
			}
		}
	}
	return dk[:keyLen]
}