package hashrng

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

	"xdx.jelly/xgcl/grand/drng/internal"

	"xdx.jelly/xgcl/sm"
	"xdx.jelly/xgcl/sm/sm3"
)

var _ = sm3.BlockSize // import sm3

// 内部状态
type internalState struct {
	v []byte
	c []byte

	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.c); i++ {
		s.c[i] = 0
	}

	s.reseedCounter = 0
}

type HashDrng struct {
	// mu                    sync.Mutex
	config                internal.DrngConfig
	hash                  sm.Hash
	df                    func(input []byte, n int) []byte
	state                 internalState
	personalizationString []byte //个性化字符串，来自应用的比特串
}

// Read implements internal.internal.
// func (hd *HashDrng) Read(p []byte) (n int, err error) {
// 	return len(p), hd.Generate(p, nil)
// }

// var SM3_df = newDF(sm.SM3)

// newDF 返回一个使用hash作为杂凑的df函数
// 例：SM3_df = newDF(sm.SM3)
func newDF(hash sm.Hash) func(input []byte, n int) []byte {
	h := hash.New()
	return func(input []byte, n int) []byte {
		out := make([]byte, 0, n+hash.Size()-1)
		len := (n + hash.Size() - 1) / hash.Size()
		counter := uint32(1)
		buf := make([]byte, 8)
		for i := 0; i < len; i++ {
			binary.BigEndian.PutUint32(buf, counter)
			binary.BigEndian.PutUint32(buf[4:], uint32(n<<3))
			h.Write(buf)
			h.Write(input)
			out = h.Sum(out)
			h.Reset()
			counter += 1
		}
		return out[:n]
	}
}

// func New(hash sm.Hash, config *internal.DrngConfig) (internal.Drng, error) {
// 	raw, _ := NewHashDrng(hash, config)
// 	pool := entropy.GetEntropyPool()
// 	d, err := raw.Instantiate([]byte(internal.PERSONALIZATIONSTRING), internal.Nonce(), pool)
// 	if err != nil {
// 		return nil, err
// 	}
// 	return &drng.DrngUtil{
// 		Drng:         d,
// 		ReseedSource: pool,
// 	}, nil
// }

func NewHashDrng(hash sm.Hash, config *internal.DrngConfig) (internal.RawDrng, error) {
	hd := &HashDrng{
		hash:   hash,
		df:     newDF(hash),
		config: *config,
	}

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

func (hd *HashDrng) OutLen() int {
	return hd.config.Outlen
}

func (hd *HashDrng) Config() *internal.DrngConfig {
	return &hd.config
}

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

	cfg := hd.config
	hd.personalizationString = personalizationString

	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 hd, err
		}
	}
	seedMaterial = append(seedMaterial, nonce...)
	seedMaterial = append(seedMaterial, personalizationString...)
	seed := hd.df(seedMaterial, cfg.Seedlen)

	//初始化内部状态变量
	hd.state.v = append(hd.state.v[:0], seed...)
	hd.state.c = append(hd.state.c[:0], hd.df(append([]byte{0}, seed...), cfg.Seedlen)...)
	hd.state.reseedCounter = 1
	hd.state.lastReseedTime = time.Now()

	// seed和seedMaterial置0
	for i := 0; i < len(seed); i++ {
		seed[i] = 0
	}
	for i := 0; i < len(seedMaterial); i++ {
		seedMaterial[i] = 0
	}
	return hd, nil
}

// Reseed 更新HashDrng的内部状态
func (hd *HashDrng) Reseed(entropyInput, additionalInput []byte) {
	// hd.mu.Lock()
	// defer hd.mu.Unlock()

	seedMaterial := make([]byte, 1, 1+len(entropyInput)+len(hd.state.v)+len(additionalInput))
	seedMaterial[0] = 1
	seedMaterial = append(seedMaterial, entropyInput...)
	seedMaterial = append(seedMaterial, hd.state.v...)
	seedMaterial = append(seedMaterial, additionalInput...)

	seed := hd.df(seedMaterial, hd.config.Seedlen)
	copy(hd.state.v, seed)
	hd.state.c = append(hd.state.c[:0], hd.df(append([]byte{0}, seed...), hd.config.Seedlen)...)
	hd.state.reseedCounter = 1
	hd.state.lastReseedTime = time.Now()

	// seed和seedMaterial置0
	for i := 0; i < len(seed); i++ {
		seed[i] = 0
	}
	for i := 0; i < len(seedMaterial); i++ {
		seedMaterial[i] = 0
	}
}

// Generate 输出函数
func (hd *HashDrng) Generate(b []byte, additionalInput []byte) error {
	// hd.mu.Lock()
	// defer hd.mu.Unlock()

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

	// 需要重播种
	if hd.state.reseedCounter > uint32(hd.config.ReseedIntervalInCounter) ||
		time.Since(hd.state.lastReseedTime) > hd.config.ReseedIntervalInTime {
		return internal.ErrNeedReseed
	}
	hash := hd.hash.New()
	if len(additionalInput) > 0 {
		hash.Write([]byte{2})
		hash.Write(hd.state.v)
		hash.Write(additionalInput)
		W := hash.Sum(nil)
		internal.AddModBytes(hd.state.v, W)
	}
	hash.Reset()
	_, _ = hash.Write(hd.state.v)
	temp := hash.Sum(nil)
	copy(b, temp)
	hash.Reset()
	hash.Write([]byte{3})
	hash.Write(hd.state.v)
	H := hash.Sum(nil)

	reseedCounterBytes := make([]byte, 4)
	binary.BigEndian.PutUint32(reseedCounterBytes, hd.state.reseedCounter)

	internal.AddModBytes(hd.state.v, H[:])
	internal.AddModBytes(hd.state.v, hd.state.c)
	internal.AddModBytes(hd.state.v, reseedCounterBytes)
	hd.state.reseedCounter++
	return nil
}