package pbkd

import (
	"errors"

	"xdx.jelly/xgcl/sm/sm4"
	"xdx.jelly/xgcl/utils/padding"
)

const (
	ENCRYPT = iota
	DECRYPT
)

type Pbeser interface {
	// EncryptBlocks
	EncryptBlocks(dst, src []byte) ([]byte, error)
	DecryptBlocks(dst, src []byte) ([]byte, error)

	EncryptWithPadding(src []byte) ([]byte, error)
	DecryptWithPadding(src []byte) ([]byte, error)

	ResetIv(iv []byte) error
}

type PbesWithSm4Cbc struct {
	key []byte
	iv  []byte
}

var errOperation = errors.New("error operation")
var errInvalidIV = errors.New("invalid iv length")

func NewPbesWithSm4Cbc(password, salt []byte, count int, iv []byte) (Pbeser, error) {
	key, err := PbkdfWithHmacSm3(password, salt, count, int64(sm4.BlockSize))
	if err != nil {
		return nil, err
	}
	if len(iv) != sm4.BlockSize {
		return nil, errInvalidIV
	}

	myIv := make([]byte, sm4.BlockSize)
	copy(myIv, iv)

	return &PbesWithSm4Cbc{
		key: key,
		iv:  myIv,
	}, nil
}

func (pbes *PbesWithSm4Cbc) EncryptBlocks(dst, src []byte) ([]byte, error) {
	return sm4.EncryptCBC(dst, pbes.iv, pbes.key, src)
}

func (pbes *PbesWithSm4Cbc) DecryptBlocks(dst, src []byte) ([]byte, error) {
	return sm4.EncryptCBC(dst, pbes.iv, pbes.key, src)
}

func (pbes *PbesWithSm4Cbc) EncryptWithPadding(src []byte) ([]byte, error) {
	pad := padding.P7.Pad(src, sm4.BlockSize)
	return sm4.EncryptCBC(nil, pbes.iv, pbes.key, pad)
}

func (pbes *PbesWithSm4Cbc) DecryptWithPadding(src []byte) ([]byte, error) {
	dst, err := sm4.DecryptCBC(nil, pbes.iv, pbes.key, src)
	if err != nil {
		return nil, err
	}

	dst, err = padding.P7.Unpad(dst, sm4.BlockSize)
	if err != nil {
		return nil, err
	}
	return dst, nil
}

func (pbes *PbesWithSm4Cbc) ResetIv(iv []byte) error {
	if len(iv) != sm4.BlockSize {
		return errInvalidIV
	}
	copy(pbes.iv, iv)
	return nil
}