package ctrrng

import (
	"crypto/cipher"
	"encoding/binary"
	"io"
	"runtime"
	"time"

	"xdx.jelly/xgcl/grand/drng/internal"
	"xdx.jelly/xgcl/internal/xor"
	"xdx.jelly/xgcl/sm/sm4"
)

type SM4BlockNewer struct{}

func (b SM4BlockNewer) BlockSize() int {
	return sm4.BlockSize
}

func (b SM4BlockNewer) New(key []byte) cipher.Block {
	if len(key) != sm4.BlockSize {
		panic("ctrrng: invalid block key length")
	}

	block, _ := sm4.NewCipher(key)
	return block
}

func NewCtrDrng(newBlock NewBlock, config *internal.DrngConfig) (internal.RawDrng, error) {
	hd := &CtrDrng{
		config:   *config,
		newBlock: newBlock,
		df:       newDF(newBlock),
	}

	// 析构时对内部状态置0
	runtime.SetFinalizer(hd, func(hd *CtrDrng) {
		hd.state.zeromize()
	})
	return hd, nil
}

type internalState struct {
	v []byte // blockLen
	k []byte // keyLen

	reseedCounter  uint32
	lastReseedTime time.Time
}

// zeromize 对内部状态置0
func (s *internalState) zeromize() {
	for i := 0; i < len(s.v); i++ {
		s.v[i] = 0
	}

	for i := 0; i < len(s.k); i++ {
		s.k[i] = 0
	}

	s.reseedCounter = 0
}

type CtrDrng struct {
	// mu sync.Mutex

	config                internal.DrngConfig
	newBlock              NewBlock
	df                    func(input []byte, n int) []byte
	state                 internalState
	personalizationString []byte //个性化字符串，来自应用的比特串
}

type NewBlock interface {
	BlockSize() int
	New(key []byte) cipher.Block
}

// newDF 返回一个使用block作为杂凑的df函数
// 例：SM3_df = newDF(sm.SM3)
func newDF(newBlock NewBlock) func(input []byte, n int) []byte {
	keylen := newBlock.BlockSize()
	outlen := newBlock.BlockSize()
	keymatial := []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 31}

	return func(input []byte, n int) []byte {
		L := len(input)

		// S = L || n || input || 0x80 || 00 such that len(S) % n = 0
		S := make([]byte, ((9+L+n-1)/n)*n)
		binary.BigEndian.PutUint32(S, uint32(L))
		binary.BigEndian.PutUint32(S[4:], uint32(n))
		copy(S[8:], input)
		S[8+L] = 0x80

		K := keymatial[:keylen]
		temp := make([]byte, 0, keylen+n)
		ivs := make([]byte, outlen+len(S)) // iv || S
		copy(ivs[4:], S)
		i := uint32(0)
		block := newBlock.New(K)
		for len(temp) < keylen+n {
			binary.BigEndian.PutUint32(ivs, i)
			i++
			temp = append(temp, cbcMac(block, ivs)...)
		}

		K = temp[:keylen]
		X := temp[keylen:]
		block = newBlock.New(K)
		temp = []byte{}
		for len(temp) < n {
			block.Encrypt(X, X)
			temp = append(temp, X...)
		}
		return temp[:n]
	}
}

// len(data) % block.BlockSize() == 0
func cbcMac(block cipher.Block, data []byte) []byte {
	blockSize := block.BlockSize()
	iv := make([]byte, blockSize)
	n := 0
	for n+blockSize <= len(data) {
		xor.XorBytes(iv, iv, data[n:])
		block.Encrypt(iv, iv)
		n += blockSize
	}
	return iv
}

func (cd *CtrDrng) keyLen() int {
	return cd.newBlock.BlockSize()
}

func (cd *CtrDrng) blockLen() int {
	return cd.newBlock.BlockSize()
}

func (cd *CtrDrng) seedLen() int {
	return cd.keyLen() + cd.OutLen()
}

func (cd *CtrDrng) OutLen() int {
	return cd.config.Outlen
}

func (cd *CtrDrng) Config() *internal.DrngConfig {
	return &cd.config
}

// Instantiate 初始化HashDrng
func (cd *CtrDrng) Instantiate(personalizationString []byte, nonce []byte, entropySource io.Reader) (internal.Drng, error) {
	// cd.mu.Lock()
	// defer cd.mu.Unlock()

	cfg := cd.config
	cd.personalizationString = personalizationString
	cd.state.k = make([]byte, cd.keyLen())
	cd.state.v = make([]byte, cd.blockLen())

	seedMaterial := make([]byte, 2*cfg.MinEntropyInputLength, int(2*cfg.MinEntropyInputLength)+len(nonce)+len(personalizationString))
	var m, n int
	var err error
	var errCount int
	for n < int(cfg.MinEntropyInputLength) {
		m, err = entropySource.Read(seedMaterial[n:])
		n += m
		if err != nil {
			errCount++
		}
		if errCount > 3 {
			return cd, err
		}
	}
	seedMaterial = append(seedMaterial, nonce...)
	seedMaterial = append(seedMaterial, personalizationString...)
	seedMaterial = cd.df(seedMaterial, cd.seedLen())
	cd.update(seedMaterial)
	return cd, nil
}

func (cd *CtrDrng) update(seedMaterial []byte) {
	seedlen := cd.seedLen()
	temp := make([]byte, seedlen)
	n := 0
	block := cd.newBlock.New(cd.state.k)
	for n+cd.blockLen() < seedlen {
		internal.IncBytes(cd.state.v)
		block.Encrypt(temp[n:n+cd.blockLen()], cd.state.v)
		n += cd.blockLen()
	}
	xor.XorBytes(temp, temp, seedMaterial)
	copy(cd.state.k, temp)
	copy(cd.state.v, temp[cd.keyLen():])
}

func (cd *CtrDrng) Reseed(entropyInput, additionalInput []byte) {
	// cd.mu.Lock()
	// defer cd.mu.Unlock()

	seedMaterial := make([]byte, len(entropyInput)+len(additionalInput))
	copy(seedMaterial, entropyInput)
	copy(seedMaterial[len(entropyInput):], additionalInput)
	seedMaterial = cd.df(seedMaterial, cd.seedLen())
	cd.update(seedMaterial)
	cd.state.reseedCounter = 1
	cd.state.lastReseedTime = time.Now()
	for i := 0; i < len(seedMaterial); i++ {
		seedMaterial[i] = 0
	}

}

func (cd *CtrDrng) Generate(b []byte, additionalInput []byte) error {
	// cd.mu.Lock()
	// defer cd.mu.Unlock()

	if len(b) > cd.config.Outlen {
		return internal.ErrReturnedBitsTooLong
	}

	// 需要重播种
	if cd.state.reseedCounter > uint32(cd.config.ReseedIntervalInCounter) ||
		time.Since(cd.state.lastReseedTime) > cd.config.ReseedIntervalInTime {
		return internal.ErrNeedReseed
	}

	if len(additionalInput) > 0 {
		additionalInput = cd.df(additionalInput, cd.seedLen())
		cd.update(additionalInput)
	} else {
		additionalInput = make([]byte, cd.seedLen())
	}
	block := cd.newBlock.New(cd.state.k)
	internal.IncBytes(cd.state.v)
	outputBlock := make([]byte, block.BlockSize())
	block.Encrypt(outputBlock, cd.state.v)
	copy(b, outputBlock)
	cd.update(additionalInput)
	cd.state.reseedCounter++
	return nil
}