package crypto

// Trinary Hamadryad hash — SWIFFT-derived lattice hash for the Gnarl scheme.
//
// Ring: Z_271[x]/(x^27 + 1), with n=27 = 3^3, p=271, m=12 input polynomials.
//
// Three output tiers share a single computation:
//
//	GnarlHash (243 bits / 31 bytes): canonical transaction identity.
//	  27 ring coefficients × 9 bits each = 243 bits.
//	  Obtained by keeping each Z_271 coefficient in full (0..270, 9 bits).
//	  Birthday bound: 2^121.5. Inversion bound: 2^243 (lattice).
//
//	GnarlMid (216 bits / 27 bytes): packet-field identity.
//	  27 coefficients × 8 bits each = 216 bits.
//	  Obtained by reducing each Z_271 coefficient mod 243 (= 3^5).
//	  Byte-aligned. Birthday bound: 2^108.
//
//	GnarlShard (54 bits / 7 bytes): epoch-scoped transaction address.
//	  27 coefficients × 2 bits each = 54 bits (27 trits).
//	  Obtained by reducing each Z_271 coefficient mod 3.
//	  Birthday bound: 2^27.
//
// Architecture (SWIFFT core, trinary variant):
//
//	Ring:       R = Z_p[x]/(x^n + 1), with n=27, p=271
//	Input:      m·n = 12·27 = 324 bits = 41 bytes (last byte partial)
//	Fixed keys: a_1 … a_m ∈ R (random, public, pre-computed in NTT form)
//
//	Steps:
//	  1. Partition input into m=12 polynomials x_i with n=27 binary coefficients
//	  2. Forward NTT: X_i = NTT(x_i) over Z_271
//	  3. Pointwise multiply: Y_i = A_i · X_i (A_i pre-stored in NTT form)
//	  4. Sum: Y = Σ Y_i
//	  5. Inverse NTT: y = INTT(Y)
//	  6. Coefficient reduction to desired tier
//	  7. Output: pack into bytes

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

// Trinary Hamadryad constants.
const (
	// GnarlM is the number of input polynomial components.
	GnarlM = 12

	// GnarlInputBits is the total input size in bits per block.
	GnarlInputBits = GnarlM * GnarlN // 324

	// gnarlInputBytes is the input block size in bytes (324 bits = 41 bytes, last byte partial).
	gnarlInputBytes = (GnarlInputBits + 7) / 8 // 41

	// GnarlHashBits is the full output size (27 × 9 = 243 bits).
	GnarlHashBits = GnarlN * 9 // 243

	// GnarlHashBytes is the full output in bytes (ceil(243/8) = 31).
	GnarlHashBytes = (GnarlHashBits + 7) / 8 // 31

	// GnarlMidBits is the mid-tier output (27 × 8 = 216 bits = 27 bytes).
	GnarlMidBits = GnarlN * 8 // 216

	// GnarlMidBytes = 27.
	GnarlMidBytes = GnarlMidBits / 8 // 27

	// GnarlShardBits is the shard output (27 × 2 = 54 bits).
	GnarlShardBits = GnarlN * 2 // 54

	// GnarlShardBytes = ceil(54/8) = 7.
	GnarlShardBytes = (GnarlShardBits + 7) / 8 // 7

	// GnarlMidR is the reduction modulus for the mid tier: 243 = 3^5.
	GnarlMidR = 243

	// GnarlShardR is the reduction modulus for the shard tier: 3.
	GnarlShardR = 3
)

// GnarlHash is a 243-bit (31-byte) canonical transaction identity.
// 27 ring coefficients in full Z_271, packed at 9 bits each.
// 5 spare bits in the last byte are zeroed.
type GnarlHash [GnarlHashBytes]byte

// GnarlMid is a 216-bit (27-byte) packet-field identity.
// 27 coefficients reduced mod 243, packed at 8 bits each. Byte-aligned.
type GnarlMid [GnarlMidBytes]byte

// GnarlShard is a 54-bit (7-byte) epoch-scoped transaction address.
// 27 coefficients reduced mod 3, packed at 2 bits each (27 trits).
// 2 spare bits in the last byte are zeroed.
type GnarlShard [GnarlShardBytes]byte

// gnarlKeys holds the m=12 fixed random polynomials a_1..a_m in NTT form.
var gnarlKeys [GnarlM][GnarlN]uint16

// gnarlKeyBasis[i][k][j] = gnarlKeys[i][j] * NTT(e_k)[j] mod 271
// where e_k is the k-th standard basis vector (1 at position k, 0 elsewhere).
// This allows computing the full NTT-domain multiply-accumulate for a binary
// input polynomial by just summing the rows for set bits — no forward NTT needed.
//
// Padded to 32 elements per row (from 27) for AVX2 alignment: 32 uint16 = 64 bytes
// = 2 YMM registers. Extra 5 elements are always zero.
const gnarlBasisPad = 32

var gnarlKeyBasis [GnarlM][GnarlN][gnarlBasisPad]uint16

func init() {
	// Ensure NTT tables are ready.
	initNTT27Tables()

	// Deterministic PRNG seeded with "gnarl-hamadryad-v1" to generate
	// the fixed key polynomials.
	seed := uint64(0)
	seedStr := []byte("gnarl-hamadryad-v1-dendrite-trinary")
	for i, b := range seedStr {
		seed ^= uint64(b) << ((uint(i) * 7) % 64)
	}

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

	for i := range GnarlM {
		var poly [GnarlN]uint16
		for j := range GnarlN {
			poly[j] = uint16(xorshift() % uint64(GnarlP))
		}
		ntt27(&poly)
		gnarlKeys[i] = poly
	}

	// Precompute basis tables: for each key i and basis vector k,
	// compute gnarlKeys[i] * NTT(e_k) pointwise.
	for i := range GnarlM {
		for k := range GnarlN {
			var ek [GnarlN]uint16
			ek[k] = 1
			ntt27(&ek) // NTT of the k-th basis vector
			for j := range GnarlN {
				gnarlKeyBasis[i][k][j] = mod271(uint32(gnarlKeys[i][j]) * uint32(ek[j]))
			}
		}
	}
}

// gnarlCompress computes one SWIFFT compression: 41 bytes → 27 coefficients in Z_271.
//
// Uses precomputed basis tables to avoid all forward NTTs. Since input polynomials
// are binary (coefficients 0 or 1), NTT(x_i) = sum of NTT(e_k) for each set bit k.
// The pointwise multiply with the key is precomputed in gnarlKeyBasis[i][k][j].
func gnarlCompress(block *[gnarlInputBytes]byte) [GnarlN]uint16 {
	// Accumulator in NTT domain (uint32, padded to 32 for AVX2 alignment).
	var acc [gnarlBasisPad]uint32

	gnarlAccumulate(&acc, &gnarlKeyBasis, block)

	// Reduce accumulated values mod 271 and convert to uint16.
	var result [GnarlN]uint16
	for j := range GnarlN {
		result[j] = mod271(acc[j])
	}

	// Inverse NTT to get coefficients.
	intt27(&result)

	return result
}

// gnarlAccumulateGeneric scans bits in block and accumulates basis table rows.
func gnarlAccumulateGeneric(acc *[gnarlBasisPad]uint32, basis *[GnarlM][GnarlN][gnarlBasisPad]uint16, block *[gnarlInputBytes]byte) {
	for i := range GnarlM {
		bitBase := i * GnarlN
		for k := range GnarlN {
			bitIdx := bitBase + k
			byteIdx := bitIdx / 8
			bitOff := uint(bitIdx % 8)
			if byteIdx < gnarlInputBytes && block[byteIdx]&(1<<bitOff) != 0 {
				row := &basis[i][k]
				for j := range gnarlBasisPad {
					acc[j] += uint32(row[j])
				}
			}
		}
	}
}

// packCoeffs243 packs 27 Z_271 coefficients (9 bits each) into 31 bytes (243 bits).
func packCoeffs243(coeffs [GnarlN]uint16) GnarlHash {
	var out GnarlHash
	bitPos := 0
	for _, c := range coeffs {
		v := c % GnarlP // ensure in [0, 270]
		for b := range 9 {
			if v&(1<<uint(b)) != 0 {
				byteIdx := bitPos / 8
				bitIdx := uint(bitPos % 8)
				out[byteIdx] |= 1 << bitIdx
			}
			bitPos++
		}
	}
	return out
}

// unpackCoeffs243 unpacks a GnarlHash back into 27 Z_271 coefficients.
func unpackCoeffs243(h GnarlHash) [GnarlN]uint16 {
	var coeffs [GnarlN]uint16
	bitPos := 0
	for i := range GnarlN {
		var v uint16
		for b := range 9 {
			byteIdx := bitPos / 8
			bitIdx := uint(bitPos % 8)
			if h[byteIdx]&(1<<bitIdx) != 0 {
				v |= 1 << uint(b)
			}
			bitPos++
		}
		coeffs[i] = v
	}
	return coeffs
}

// packCoeffs216 packs 27 coefficients reduced mod 243 (8 bits each) into 27 bytes.
func packCoeffs216(coeffs [GnarlN]uint16) GnarlMid {
	var out GnarlMid
	for i, c := range coeffs {
		out[i] = byte(c % GnarlMidR)
	}
	return out
}

// packCoeffs54 packs 27 coefficients reduced mod 3 (2 bits each) into 7 bytes.
func packCoeffs54(coeffs [GnarlN]uint16) GnarlShard {
	var out GnarlShard
	bitPos := 0
	for _, c := range coeffs {
		v := c % GnarlShardR // 0, 1, or 2
		for b := range 2 {
			if v&(1<<uint(b)) != 0 {
				byteIdx := bitPos / 8
				bitIdx := uint(bitPos % 8)
				out[byteIdx] |= 1 << bitIdx
			}
			bitPos++
		}
	}
	return out
}

// gnarlHashCore computes the trinary Hamadryad hash of an arbitrary-length message,
// returning the raw Z_271 coefficients. All output tiers derive from this.
func gnarlHashCore(msg []byte) [GnarlN]uint16 {
	var chain [GnarlN]uint16

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

	// Pad to next multiple of gnarlInputBytes, leaving 8 bytes for length.
	for (len(padded)+8)%gnarlInputBytes != 0 {
		padded = append(padded, 0)
	}

	// Append message length in bits as 8-byte LE.
	var lenBuf [8]byte
	binary.LittleEndian.PutUint64(lenBuf[:], uint64(len(msg))*8)
	padded = append(padded, lenBuf[:]...)

	// Process each block.
	for off := 0; off < len(padded); off += gnarlInputBytes {
		var block [gnarlInputBytes]byte
		copy(block[:], padded[off:off+gnarlInputBytes])

		result := gnarlCompress(&block)

		// Chain: coefficient-wise addition mod 271.
		for j := range GnarlN {
			chain[j] = mod271(uint32(chain[j]) + uint32(result[j]))
		}
	}

	return chain
}

// GHash computes the 243-bit GnarlHash of an arbitrary-length message.
func GHash(msg []byte) GnarlHash {
	coeffs := gnarlHashCore(msg)
	return packCoeffs243(coeffs)
}

// GMid computes the 216-bit (27-byte) GnarlMid of an arbitrary-length message.
func GMid(msg []byte) GnarlMid {
	coeffs := gnarlHashCore(msg)
	return packCoeffs216(coeffs)
}

// GShard computes the 54-bit GnarlShard of an arbitrary-length message.
func GShard(msg []byte) GnarlShard {
	coeffs := gnarlHashCore(msg)
	return packCoeffs54(coeffs)
}

// Mid extracts the GnarlMid tier from a full GnarlHash.
func (h GnarlHash) Mid() GnarlMid {
	coeffs := unpackCoeffs243(h)
	return packCoeffs216(coeffs)
}

// Shard extracts the GnarlShard tier from a full GnarlHash.
func (h GnarlHash) Shard() GnarlShard {
	coeffs := unpackCoeffs243(h)
	return packCoeffs54(coeffs)
}

// IsZero reports whether the GnarlHash is all zeros.
func (h GnarlHash) IsZero() bool {
	for _, b := range h {
		if b != 0 {
			return false
		}
	}
	return true
}

// IsZero reports whether the GnarlMid is all zeros.
func (m GnarlMid) IsZero() bool {
	for _, b := range m {
		if b != 0 {
			return false
		}
	}
	return true
}

// IsZero reports whether the GnarlShard is all zeros.
func (s GnarlShard) IsZero() bool {
	for _, b := range s {
		if b != 0 {
			return false
		}
	}
	return true
}

// Sum returns a new GnarlHash that is the coefficient-wise sum (mod 271)
// of two GnarlHash values.
func (h GnarlHash) Sum(other GnarlHash) GnarlHash {
	a := unpackCoeffs243(h)
	b := unpackCoeffs243(other)
	var result [GnarlN]uint16
	for i := range GnarlN {
		result[i] = mod271(uint32(a[i]) + uint32(b[i]))
	}
	return packCoeffs243(result)
}

// GnarlHashValue computes the GnarlHash of an arbitrary Go value.
func GnarlHashValue(v any) GnarlHash {
	data := []byte(fmt.Sprintf("%v", v))
	return GHash(data)
}

// RandomGnarlHash generates a cryptographically random GnarlHash value.
func RandomGnarlHash() GnarlHash {
	var h GnarlHash
	if _, err := rand.Read(h[:]); err != nil {
		panic(fmt.Sprintf("crypto/rand: %v", err))
	}
	// Zero the 5 spare bits in the last byte.
	h[GnarlHashBytes-1] &= 0x07
	return h
}