// Ring-SIS lattice hash interface.
//
// The Short Integer Solution (SIS) problem over polynomial rings:
// given random a_1, ..., a_m ∈ R_q = Z_q[x]/(x^n + 1), find short
// x_1, ..., x_m ∈ R_q such that Σ a_i · x_i = 0 (mod q).
//
// This is provably hard, reducing to worst-case SVP on ideal lattices.
//
// A SWIFFT-family hash is the compression function:
//   f_A(x) = Σ a_i · x_i   (NTT-accelerated)
//
// where x_i have binary (0/1) or short integer coefficients.
//
// Properties:
//   - Collision resistance: finding x ≠ x' with f(x) = f(x') ↔ solving Ring-SIS
//   - Additive homomorphism: f(x + y) = f(x) + f(y) when x+y stays short
//   - Universal hashing: for random A, uniform over the output ring
//
// This package provides SISParams and SISHasher to formalize the pattern
// used by both hamadryad (n=64, p=257) and gnarl (n=27, p=271).

package ring

import (
	"crypto/rand"
	"encoding/binary"
	"io"
)

// SISParams defines a Ring-SIS hash family instance.
type SISParams struct {
	// Ring dimension and modulus.
	N int    // polynomial degree (power of 2 for hamadryad, power of 3 for gnarl)
	Q uint32 // coefficient modulus (prime, q ≡ 1 mod 2n for NTT)

	// M is the number of key polynomials (input components).
	// Input is M binary polynomials of degree N, so input size = M·N bits.
	M int

	// ReductionMod is the output reduction modulus.
	// Each coefficient is reduced mod ReductionMod for the final hash.
	// For hamadryad: 128 (7 bits). For gnarl: 271/243/3 depending on tier.
	ReductionMod uint32

	// BitsPerCoeff is the number of bits per output coefficient.
	BitsPerCoeff int

	// InputBits is the total input size (M * N).
	InputBits int

	// OutputBits is the total output size (N * BitsPerCoeff).
	OutputBits int
}

// HamadryadSISParams returns the SIS parameters for the hamadryad hash.
//
//	n=64, p=257, m=16, output reduced mod 128 (7 bits/coeff) = 448 bits.
func HamadryadSISParams() SISParams {
	return SISParams{
		N:            64,
		Q:            257,
		M:            16,
		ReductionMod: 128,
		BitsPerCoeff: 7,
		InputBits:    16 * 64, // 1024
		OutputBits:   64 * 7,  // 448
	}
}

// GnarlSISParams returns the SIS parameters for the gnarl hash (full tier).
//
//	n=27, p=271, m=12, output unreduced (9 bits/coeff) = 243 bits.
func GnarlSISParams() SISParams {
	return SISParams{
		N:            27,
		Q:            271,
		M:            12,
		ReductionMod: 271, // full range
		BitsPerCoeff: 9,
		InputBits:    12 * 27, // 324
		OutputBits:   27 * 9,  // 243
	}
}

// SISHasher computes Ring-SIS hashes using the SWIFFT construction.
// Thread-safe after initialization (all key material is read-only).
type SISHasher struct {
	params SISParams
	keys   []*Poly // M key polynomials in NTT form

	// nttFwd and nttInv are NTT functions for the specific ring.
	// For standard power-of-2 dimensions, these use the ring package NTT.
	// For non-power-of-2 (gnarl's n=27), these would need the specialized NTT.
	ringParams Params
}

// NewSISHasher creates a SIS hasher with deterministic key generation.
// The seed string determines the key polynomials.
// For interoperability with existing hamadryad, use "hamadryad-swifft-v1-dendrite-kismet".
func NewSISHasher(sp SISParams, seed string) *SISHasher {
	// Compute ring params. For SIS, we need NTT support,
	// which requires q ≡ 1 (mod 2n) and a primitive root.
	rp := Params{
		N: sp.N,
		Q: sp.Q,
	}

	// Find a primitive 2N-th root of unity for the NTT.
	// For known parameter sets, use precomputed values.
	switch {
	case sp.N == 64 && sp.Q == 257:
		rp.RootOfUnity = 9 // psi=9, primitive 128th root of unity mod 257
		rp.MontR = 1 << 16
		rp.QInv = qinv(257, 16)
	case sp.N == 27 && sp.Q == 271:
		// Gnarl uses a specialized 27-point NTT (not power-of-2).
		// The ring package NTT won't work directly for this.
		// For now, mark as needing external NTT.
		rp.RootOfUnity = 0 // signal: needs external NTT
	default:
		// For other parameter sets, try to find a root.
		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,
	}

	// Generate key polynomials deterministically from seed.
	h.keys = h.genKeys(seed)

	return h
}

// genKeys generates M key polynomials from a deterministic seed.
// Uses the same xorshift PRNG as the existing hamadryad implementation.
func (h *SISHasher) genKeys(seed string) []*Poly {
	// Derive PRNG state from seed string.
	state := uint64(0)
	for i, b := range []byte(seed) {
		state ^= uint64(b) << ((uint(i) * 7) % 64)
	}

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

	keys := make([]*Poly, h.params.M)
	for i := range keys {
		poly := New(h.ringParams)
		for j := range poly.Coeffs {
			poly.Coeffs[j] = uint32(xorshift() % uint64(h.params.Q))
		}
		// Store in NTT form for fast multiplication.
		if h.ringParams.RootOfUnity != 0 {
			NTT(poly)
		}
		keys[i] = poly
	}
	return keys
}

// Compress computes one SWIFFT compression: M·N bits → N coefficients in Z_q.
// The input block must be exactly InputBits/8 bytes.
// Returns the raw coefficients (before output reduction).
func (h *SISHasher) Compress(block []byte) *Poly {
	sp := h.params
	acc := New(h.ringParams)

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

		// NTT → pointwise multiply with key → accumulate.
		NTT(poly)
		MulAdd(acc, poly, h.keys[i])
	}

	acc.isNTT = true

	// Inverse NTT to get coefficients.
	INTT(acc)

	return acc
}

// Hash computes the Ring-SIS hash of an arbitrary-length message.
// Uses Merkle-Damgård chaining: each block is compressed, and the
// result is added (coefficient-wise mod q) to the chain.
func (h *SISHasher) Hash(msg []byte) *Poly {
	sp := h.params
	blockBytes := sp.InputBits / 8
	if sp.InputBits%8 != 0 {
		blockBytes++
	}

	var chain *Poly

	// Pad message: 0x80 || zeros || 8-byte LE length.
	padded := make([]byte, len(msg))
	copy(padded, msg)
	padded = append(padded, 0x80)

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

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

	// Process blocks.
	for off := 0; off < 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
}

// ReduceAndPack reduces coefficients mod ReductionMod and packs them
// into a byte slice at BitsPerCoeff bits each.
func (h *SISHasher) ReduceAndPack(coeffs *Poly) []byte {
	sp := h.params
	totalBits := sp.N * sp.BitsPerCoeff
	out := make([]byte, (totalBits+7)/8)

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

// HashBytes computes the full hash (compress + reduce + pack) of a message.
func (h *SISHasher) HashBytes(msg []byte) []byte {
	coeffs := h.Hash(msg)
	return h.ReduceAndPack(coeffs)
}

// SumCoeffs computes the coefficient-wise sum of two hash outputs (mod q).
// This is the additive homomorphism: Hash(a ⊕ b) = Hash(a) + Hash(b)
// when the combined input stays short.
func (h *SISHasher) SumCoeffs(a, b *Poly) *Poly {
	return Add(a, b)
}

// Params returns the SIS parameters.
func (h *SISHasher) Params() SISParams {
	return h.params
}

// RingParams returns the underlying ring parameters.
func (h *SISHasher) RingParams() Params {
	return h.ringParams
}

// Keys returns the key polynomials (in NTT form).
func (h *SISHasher) Keys() []*Poly {
	return h.keys
}

// findPrimRoot finds a primitive k-th root of unity mod q.
// Returns 0 if none exists (i.e., k does not divide q-1).
func findPrimRoot(q, k uint32) uint32 {
	if (q-1)%k != 0 {
		return 0
	}

	// Find a generator of Z_q* by trial.
	// For small primes, just try candidates.
	for g := uint32(2); g < q; g++ {
		// Check if g is a primitive root mod q.
		if powModU32(g, (q-1)/2, q) == 1 {
			continue // not a generator
		}
		// g^((q-1)/k) is a primitive k-th root.
		root := powModU32(g, (q-1)/k, q)
		if root != 1 {
			return root
		}
	}
	return 0
}

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

// NewSISHasherRandom creates a SIS hasher with randomly generated keys.
// Used when you don't need interoperability with a fixed key set.
func NewSISHasherRandom(sp SISParams) *SISHasher {
	return NewSISHasherRandomFrom(sp, rand.Reader)
}

// NewSISHasherRandomFrom creates a SIS hasher with keys from the given source.
func NewSISHasherRandomFrom(sp SISParams, rng io.Reader) *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 := make([]*Poly, sp.M)
	for i := range keys {
		keys[i] = UniformPolyFrom(rp, rng)
		NTT(keys[i])
	}
	h.keys = keys

	return h
}