package gnarlring

import "git.smesh.lol/gnarl-hamadryad/crypto"

const (
	N            = crypto.GnarlN // 27
	Q            = crypto.GnarlP // 271
	bitsPerCoeff = 9             // ceil(log2(271))
	PolyBytes    = 31            // ceil(27*9/8)
)

// Poly27 is a polynomial in Z_271[x]/(x^27+1). Stored as 27 uint16
// coefficients in [0, 271). Can be in coefficient form or NTT form.
type Poly27 struct {
	Coeffs [N]uint16
	isNTT  bool
}

// NewPoly27 returns the zero polynomial in coefficient form.
func NewPoly27() *Poly27 {
	return &Poly27{}
}

// Clone returns a deep copy.
func (p *Poly27) Clone() *Poly27 {
	c := &Poly27{isNTT: p.isNTT}
	c.Coeffs = p.Coeffs
	return c
}

// Set copies src into p.
func (p *Poly27) Set(src *Poly27) {
	p.Coeffs = src.Coeffs
	p.isNTT = src.isNTT
}

// Zero sets all coefficients to zero.
func (p *Poly27) Zero() {
	p.Coeffs = [N]uint16{}
	p.isNTT = false
}

// IsNTT reports whether the polynomial is in NTT form.
func (p *Poly27) IsNTT() bool { return p.isNTT }

// --- Coefficient arithmetic (all mod 271) ---

// Add returns a + b (coefficient-wise mod 271). a and b must be in the same form.
func Add(a, b *Poly27) *Poly27 {
	c := &Poly27{isNTT: a.isNTT}
	for i := range a.Coeffs {
		c.Coeffs[i] = addMod(a.Coeffs[i], b.Coeffs[i])
	}
	return c
}

// Sub returns a - b (coefficient-wise mod 271).
func Sub(a, b *Poly27) *Poly27 {
	c := &Poly27{isNTT: a.isNTT}
	for i := range a.Coeffs {
		c.Coeffs[i] = subMod(a.Coeffs[i], b.Coeffs[i])
	}
	return c
}

// Neg returns -a (coefficient-wise mod 271).
func Neg(a *Poly27) *Poly27 {
	c := &Poly27{isNTT: a.isNTT}
	for i, v := range a.Coeffs {
		if v != 0 {
			c.Coeffs[i] = Q - v
		}
	}
	return c
}

// ScalarMul returns s * a (coefficient-wise mod 271).
func ScalarMul(a *Poly27, s uint16) *Poly27 {
	c := &Poly27{isNTT: a.isNTT}
	ss := uint32(s % Q)
	for i, v := range a.Coeffs {
		c.Coeffs[i] = crypto.Mod271(uint32(v) * ss)
	}
	return c
}

func addMod(a, b uint16) uint16 {
	s := uint32(a) + uint32(b)
	if s >= Q {
		s -= Q
	}
	return uint16(s)
}

func subMod(a, b uint16) uint16 {
	if a >= b {
		return a - b
	}
	return Q - b + a
}

// --- NTT operations ---

// NTT computes the forward NTT in-place. Delegates to crypto.NTT27.
func (p *Poly27) NTT() {
	if p.isNTT {
		return
	}
	crypto.NTT27(&p.Coeffs)
	p.isNTT = true
}

// INTT computes the inverse NTT in-place. Delegates to crypto.INTT27.
func (p *Poly27) INTT() {
	if !p.isNTT {
		return
	}
	crypto.INTT27(&p.Coeffs)
	p.isNTT = false
}

// --- Multiplication ---

// MulPointwise returns a · b pointwise in the NTT domain. Both must be in NTT form.
func MulPointwise(a, b *Poly27) *Poly27 {
	c := &Poly27{isNTT: true}
	for i := range a.Coeffs {
		c.Coeffs[i] = crypto.Mod271(uint32(a.Coeffs[i]) * uint32(b.Coeffs[i]))
	}
	return c
}

// Mul returns a * b in the ring Z_271[x]/(x^27+1). Uses NTT pipeline: converts
// both inputs to NTT, multiplies pointwise, converts back. Returns coefficient form.
// Inputs are unmodified.
func Mul(a, b *Poly27) *Poly27 {
	ca := a.Clone()
	cb := b.Clone()
	ca.NTT()
	cb.NTT()
	c := MulPointwise(ca, cb)
	c.INTT()
	return c
}

// --- Ring inversion ---

// Inverse returns a^{-1} in the ring Z_271[x]/(x^27+1). Uses NTT domain
// pointwise inversion via Fermat's little theorem. Returns nil if any
// NTT slot is zero (a is not invertible). a must be in coefficient form.
func Inverse(a *Poly27) *Poly27 {
	c := a.Clone()
	c.NTT()
	for i, v := range c.Coeffs {
		if v == 0 {
			return nil
		}
		c.Coeffs[i] = crypto.InvMod271(v)
	}
	c.INTT()
	return c
}

// --- Norm computation ---

// Norm returns the centered infinity norm: max absolute coefficient value
// when coefficients are represented as signed integers centered at Q/2.
func Norm(a *Poly27) uint16 {
	var m uint16
	half := uint16(Q / 2)
	for _, v := range a.Coeffs {
		var abs uint16
		if v > half {
			abs = Q - v
		} else {
			abs = v
		}
		if abs > m {
			m = abs
		}
	}
	return m
}

// NormSq returns the squared Euclidean norm (sum of squared centered
// coefficients as integers).
func NormSq(a *Poly27) uint64 {
	var sum uint64
	half := uint64(Q / 2)
	for _, v := range a.Coeffs {
		vc := uint64(v)
		var abs uint64
		if vc > half {
			abs = uint64(Q) - vc
		} else {
			abs = vc
		}
		sum += abs * abs
	}
	return sum
}

// Dot returns the Euclidean dot product of a and b, using centered
// (signed) coefficient values. The result is a signed integer whose
// magnitude is bounded by n · Q².
func Dot(a, b *Poly27) int64 {
	var sum int64
	half := int64(Q / 2)
	for i := range a.Coeffs {
		ai := int64(a.Coeffs[i])
		bi := int64(b.Coeffs[i])
		if ai > half {
			ai -= int64(Q)
		}
		if bi > half {
			bi -= int64(Q)
		}
		sum += ai * bi
	}
	return sum
}

// --- Comparison ---

// Equal reports whether two polynomials have identical coefficients.
func Equal(a, b *Poly27) bool {
	return a.Coeffs == b.Coeffs && a.isNTT == b.isNTT
}

// IsZero reports whether all coefficients are zero.
func IsZero(a *Poly27) bool {
	for _, v := range a.Coeffs {
		if v != 0 {
			return false
		}
	}
	return true
}

// --- Serialization ---

// Serialize packs polynomial coefficients at the given bit width using LE
// bitwise packing. If signed is true, coefficients are stored in two's
// complement (caller must ensure |centered_coeff| < 2^(bitsPerCoeff-1)).
func Serialize(a *Poly27, bitsPerCoeff int, signed bool) []byte {
	totalBits := bitsPerCoeff * N
	out := make([]byte, (totalBits+7)/8)

	mask := uint32((1 << bitsPerCoeff) - 1)
	bitPos := 0

	for _, c := range a.Coeffs {
		var v uint32
		if signed {
			half := uint16(Q / 2)
			if c > half {
				abs := uint32(Q - c)
				v = ((^abs) + 1) & mask
			} else {
				v = uint32(c) & mask
			}
		} else {
			v = uint32(c) & mask
		}
		for b := 0; b < bitsPerCoeff; b++ {
			if v&(1<<uint(b)) != 0 {
				out[bitPos/8] |= 1 << uint(bitPos%8)
			}
			bitPos++
		}
	}
	return out
}

// Deserialize unpacks a polynomial serialized by Serialize. Returns nil if
// the data buffer is too short.
func Deserialize(data []byte, bitsPerCoeff int, signed bool) *Poly27 {
	totalBits := bitsPerCoeff * N
	if len(data) < (totalBits+7)/8 {
		return nil
	}

	p := NewPoly27()
	mask := uint32((1 << bitsPerCoeff) - 1)
	var signBit uint32
	if signed {
		signBit = 1 << (bitsPerCoeff - 1)
	}
	bitPos := 0

	for i := 0; i < N; i++ {
		var v uint32
		for b := 0; b < bitsPerCoeff; b++ {
			if bitPos/8 < len(data) && data[bitPos/8]&(1<<uint(bitPos%8)) != 0 {
				v |= 1 << uint(b)
			}
			bitPos++
		}
		v &= mask

		if signed && v >= signBit {
			abs := ((^v) + 1) & mask
			p.Coeffs[i] = Q - uint16(abs)
			if p.Coeffs[i] >= Q {
				p.Coeffs[i] = 0
			}
		} else {
			p.Coeffs[i] = uint16(v)
		}
	}
	return p
}

// MarshalBinary serializes at 9 bits/coeff unsigned. Returns exactly
// PolyBytes (31) bytes.
func (p *Poly27) MarshalBinary() []byte {
	return Serialize(p, bitsPerCoeff, false)
}

// UnmarshalBinary deserializes a polynomial encoded by MarshalBinary.
// Returns an error if the data is too short.
func UnmarshalBinary(data []byte) (*Poly27, error) {
	p := Deserialize(data, bitsPerCoeff, false)
	if p == nil {
		return nil, errShortData
	}
	return p, nil
}

// errShortData is a sentinel for deserialization failures.
var errShortData = errBytes("gnarlring: data too short")

type errBytes string

func (e errBytes) Error() string { return string(e) }