package sm234

import (
	"encoding/binary"
	"hash"
	"math/bits"
)

/*
// hash.Hash
type Hash interface {
	io.Writer
	Sum(b []byte) []byte
	Reset()
	Size() int
	BlockSize() int
}
*/

type sm3ctx struct {
	digest [8]uint32
	block  [64]byte

	offs int // block 中等待处理的字节数.
	nblk int // 已经处理过的 block 个数.
}

func NewSM3Hash() hash.Hash {
	var nctx sm3ctx
	nctx.Reset()
	return &nctx
}

func (ctx *sm3ctx) Size() int {
	return 32
}

func (ctx *sm3ctx) BlockSize() int {
	return 64
}

func (ctx *sm3ctx) Reset() {
	ctx.digest[0] = 0x7380166F
	ctx.digest[1] = 0x4914B2B9
	ctx.digest[2] = 0x172442D7
	ctx.digest[3] = 0xDA8A0600
	ctx.digest[4] = 0xA96F30BC
	ctx.digest[5] = 0x163138AA
	ctx.digest[6] = 0xE38DEE4D
	ctx.digest[7] = 0xB0FB0E4E

	ctx.offs = 0
	ctx.nblk = 0
}

// block
func (ctx *sm3ctx) round(blk []byte) {
	var A, B, C, D, E, F, G, H uint32
	var W [68]uint32
	var SS1, SS2, TT1, TT2 uint32
	//
	A = ctx.digest[0]
	B = ctx.digest[1]
	C = ctx.digest[2]
	D = ctx.digest[3]
	E = ctx.digest[4]
	F = ctx.digest[5]
	G = ctx.digest[6]
	H = ctx.digest[7]

	// 64 bytes per BLOCK.
	for i := 0; i < 16; i++ {
		W[i] = binary.BigEndian.Uint32(blk[i*4:])
	}

	for i := 16; i < 68; i++ {
		// P1(x) = ((x) ^ ROTL((x),15) ^ ROTL((x),23))
		temp := W[i-16] ^ W[i-9] ^ bits.RotateLeft32(W[i-3], 15)
		temp = temp ^ bits.RotateLeft32(temp, 15) ^ bits.RotateLeft32(temp, 23)
		W[i] = temp ^ bits.RotateLeft32(W[i-13], 7) ^ W[i-6]
	}

	//
	for i := 0; i < 16; i++ {
		KK := bits.RotateLeft32(A, 12) + E + bits.RotateLeft32(0x79cc4519, i)
		SS1 = bits.RotateLeft32(KK, 7)
		SS2 = SS1 ^ bits.RotateLeft32(A, 12)

		TT1 = (A ^ B ^ C) + D + SS2 + (W[i] ^ W[i+4])
		TT2 = (E ^ F ^ G) + H + SS1 + W[i]

		D = C
		C = bits.RotateLeft32(B, 9)
		B = A
		A = TT1
		H = G
		G = bits.RotateLeft32(F, 19)
		F = E
		E = TT2 ^ bits.RotateLeft32(TT2, 9) ^ bits.RotateLeft32(TT2, 17)
	}

	for i := 16; i < 64; i++ {
		KK := bits.RotateLeft32(A, 12) + E + bits.RotateLeft32(0x7a879d8a, i)
		SS1 = bits.RotateLeft32(KK, 7)
		SS2 = SS1 ^ bits.RotateLeft32(A, 12)

		TT1 = ((A & B) | (A & C) | (B & C)) + D + SS2 + (W[i] ^ W[i+4])
		TT2 = (((F ^ G) & E) ^ G) + H + SS1 + W[i]

		D = C
		C = bits.RotateLeft32(B, 9)
		B = A
		A = TT1
		H = G
		G = bits.RotateLeft32(F, 19)
		F = E
		E = TT2 ^ bits.RotateLeft32(TT2, 9) ^ bits.RotateLeft32(TT2, 17)
	}

	//
	ctx.digest[0] ^= A
	ctx.digest[1] ^= B
	ctx.digest[2] ^= C
	ctx.digest[3] ^= D
	ctx.digest[4] ^= E
	ctx.digest[5] ^= F
	ctx.digest[6] ^= G
	ctx.digest[7] ^= H

}

// adds more data to the running hash.
// It never returns an error.
func (ctx *sm3ctx) Write(p []byte) (n int, err error) {

	var tsiz int

	for i := 0; i < len(p); i += tsiz {

		tsiz = copy(ctx.block[ctx.offs:], p[i:])

		ctx.offs += tsiz
		if ctx.offs < 64 {
			break
		}

		// block is full..
		ctx.round(ctx.block[:])
		ctx.offs = 0
		ctx.nblk += 1
	}

	return len(p), nil
}

// Sum appends the current hash to b and returns the resulting slice.
// It does not change the underlying hash state.
func (ctx *sm3ctx) Sum(b []byte) []byte {

	// copy a new context
	nctx := *ctx
	spad := make([]byte, 64)

	// padding tail, add bytes count...
	nctx.block[nctx.offs] = 0x80
	copy(spad[:], nctx.block[0:nctx.offs+1])

	//fmt.Printf("tail: %v\n", nctx.offs)
	//fmt.Println(hex.EncodeToString(spad))

	if nctx.offs > (64 - 9) {
		nctx.round(spad[:])

		// zero 56 bytes.
		for i := 0; i < 56; i++ {
			spad[i] = 0
		}
	}

	//
	binary.BigEndian.PutUint32(spad[56:], uint32(nctx.nblk)>>23)
	binary.BigEndian.PutUint32(spad[60:], (uint32(nctx.nblk)<<9)+(uint32(nctx.offs)<<3))
	nctx.round(spad[:])

	//
	tofs := len(b)
	rult := make([]byte, tofs+32)
	if tofs > 0 {
		copy(rult[0:], b[:])
	}

	binary.BigEndian.PutUint32(rult[tofs:], nctx.digest[0])
	binary.BigEndian.PutUint32(rult[tofs+4:], nctx.digest[1])
	binary.BigEndian.PutUint32(rult[tofs+8:], nctx.digest[2])
	binary.BigEndian.PutUint32(rult[tofs+12:], nctx.digest[3])
	binary.BigEndian.PutUint32(rult[tofs+16:], nctx.digest[4])
	binary.BigEndian.PutUint32(rult[tofs+20:], nctx.digest[5])
	binary.BigEndian.PutUint32(rult[tofs+24:], nctx.digest[6])
	binary.BigEndian.PutUint32(rult[tofs+28:], nctx.digest[7])
	return rult
}