// Ring-SIS lattice hash interface (SWIFFT construction).
package ring

import (
	"crypto/rand"
	"io"
)

type SISParams struct {
	N            int32
	Q            uint32
	M            int32
	ReductionMod uint32
	BitsPerCoeff int32
	InputBits    int32
	OutputBits   int32
}

func HamadryadSISParams() (sp SISParams) {
	return SISParams{
		N:            64,
		Q:            257,
		M:            16,
		ReductionMod: 128,
		BitsPerCoeff: 7,
		InputBits:    16 * 64,
		OutputBits:   64 * 7,
	}
}

func GnarlSISParams() (sp SISParams) {
	return SISParams{
		N:            27,
		Q:            271,
		M:            12,
		ReductionMod: 271,
		BitsPerCoeff: 9,
		InputBits:    12 * 27,
		OutputBits:   27 * 9,
	}
}

type SISHasher struct {
	params     SISParams
	keys       []*Poly
	ringParams Params
}

func NewSISHasher(sp SISParams, seed string) (h *SISHasher) {
	rp := Params{
		N: sp.N,
		Q: sp.Q,
	}

	switch {
	case sp.N == 64 && sp.Q == 257:
		rp.RootOfUnity = 9
		rp.MontR = 1 << 16
		rp.QInv = qinv(257, 16)
	case sp.N == 27 && sp.Q == 271:
		rp.RootOfUnity = 0
	default:
		rp.RootOfUnity = findPrimRoot(sp.Q, uint32(2*sp.N))
		if sp.Q < (1 << 16) {
			rp.MontR = 1 << 16
			rp.QInv = qinv(sp.Q, 16)
		} else {
			rp.MontR = 1 << 32
			rp.QInv = qinv(sp.Q, 32)
		}
	}

	h = &SISHasher{
		params:     sp,
		ringParams: rp,
	}

	h.keys = h.genKeys(seed)
	return h
}

func (h *SISHasher) genKeys(seed string) (keys []*Poly) {
	state := uint64(0)
	seedBytes := []byte(seed)
	for i := int32(0); i < int32(len(seedBytes)); i++ {
		state ^= uint64(seedBytes[i]) << ((uint32(i) * 7) % 64)
	}

	xorshift := func() (v uint64) {
		state ^= state << 13
		state ^= state >> 7
		state ^= state << 17
		return state
	}

	keys = []*Poly{:h.params.M}
	for i := int32(0); i < h.params.M; i++ {
		poly := New(h.ringParams)
		for j := int32(0); j < int32(len(poly.Coeffs)); j++ {
			poly.Coeffs[j] = uint32(xorshift() % uint64(h.params.Q))
		}
		if h.ringParams.RootOfUnity != 0 {
			NTT(poly)
		}
		keys[i] = poly
	}
	return keys
}

func (h *SISHasher) Compress(block []byte) (acc *Poly) {
	sp := h.params
	acc = New(h.ringParams)

	for i := int32(0); i < sp.M; i++ {
		poly := New(h.ringParams)
		bitBase := i * sp.N
		for j := int32(0); j < sp.N; j++ {
			bitIdx := bitBase + j
			byteIdx := bitIdx / 8
			bitOff := uint32(bitIdx % 8)
			if byteIdx < int32(len(block)) && block[byteIdx]&(1<<bitOff) != 0 {
				poly.Coeffs[j] = 1
			}
		}

		NTT(poly)
		MulAdd(acc, poly, h.keys[i])
	}

	acc.isNTT = true
	INTT(acc)
	return acc
}

func (h *SISHasher) Hash(msg []byte) (chain *Poly) {
	sp := h.params
	blockBytes := sp.InputBits / 8
	if sp.InputBits%8 != 0 {
		blockBytes++
	}

	padded := []byte{:int32(len(msg))}
	copy(padded, msg)
	padded = append(padded, 0x80)

	for (int32(len(padded))+8)%blockBytes != 0 {
		padded = append(padded, 0)
	}

	var lenBuf [8]byte
	putLeU64(lenBuf[:], uint64(len(msg))*8)
	padded = append(padded, lenBuf[:]...)

	for off := int32(0); off < int32(len(padded)); off += blockBytes {
		result := h.Compress(padded[off : off+blockBytes])

		if chain == nil {
			chain = result
		} else {
			chain = Add(chain, result)
		}
	}

	if chain == nil {
		chain = New(h.ringParams)
	}
	return chain
}

func (h *SISHasher) ReduceAndPack(coeffs *Poly) (out []byte) {
	sp := h.params
	totalBits := sp.N * sp.BitsPerCoeff
	out = []byte{:(totalBits + 7) / 8}

	bitPos := int32(0)
	for i := int32(0); i < sp.N; i++ {
		v := coeffs.Coeffs[i] % sp.ReductionMod
		for b := int32(0); b < sp.BitsPerCoeff; b++ {
			if v&(1<<uint32(b)) != 0 {
				out[bitPos/8] |= 1 << uint32(bitPos%8)
			}
			bitPos++
		}
	}
	return out
}

func (h *SISHasher) HashBytes(msg []byte) (out []byte) {
	coeffs := h.Hash(msg)
	return h.ReduceAndPack(coeffs)
}

func (h *SISHasher) SumCoeffs(a, b *Poly) (c *Poly) {
	return Add(a, b)
}

func (h *SISHasher) SISParams() (sp SISParams) {
	return h.params
}

func (h *SISHasher) RingParams() (rp Params) {
	return h.ringParams
}

func (h *SISHasher) Keys() (keys []*Poly) {
	return h.keys
}

func findPrimRoot(q, k uint32) (root uint32) {
	if (q-1)%k != 0 {
		return 0
	}

	for g := uint32(2); g < q; g++ {
		if powModU32(g, (q-1)/2, q) == 1 {
			continue
		}
		root = powModU32(g, (q-1)/k, q)
		if root != 1 {
			return root
		}
	}
	return 0
}

func powModU32(base, exp, m uint32) (result uint32) {
	r := uint64(1)
	b := uint64(base) % uint64(m)
	for e := exp; e > 0; e >>= 1 {
		if e&1 == 1 {
			r = (r * b) % uint64(m)
		}
		b = (b * b) % uint64(m)
	}
	return uint32(r)
}

func NewSISHasherRandom(sp SISParams) (h *SISHasher) {
	return NewSISHasherRandomFrom(sp, rand.Reader)
}

func NewSISHasherRandomFrom(sp SISParams, rng io.Reader) (h *SISHasher) {
	rp := Params{
		N: sp.N,
		Q: sp.Q,
	}

	switch {
	case sp.N == 64 && sp.Q == 257:
		rp.RootOfUnity = 9
		rp.MontR = 1 << 16
		rp.QInv = qinv(257, 16)
	default:
		rp.RootOfUnity = findPrimRoot(sp.Q, uint32(2*sp.N))
		if sp.Q < (1 << 16) {
			rp.MontR = 1 << 16
			rp.QInv = qinv(sp.Q, 16)
		} else {
			rp.MontR = 1 << 32
			rp.QInv = qinv(sp.Q, 32)
		}
	}

	h = &SISHasher{
		params:     sp,
		ringParams: rp,
	}

	keys := []*Poly{:sp.M}
	for i := int32(0); i < sp.M; i++ {
		keys[i] = UniformPolyFrom(rp, rng)
		NTT(keys[i])
	}
	h.keys = keys

	return h
}