package ring

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

// UniformPoly samples a polynomial with coefficients uniformly random in [0, Q).
func UniformPoly(p Params) *Poly {
	return UniformPolyFrom(p, rand.Reader)
}

// UniformPolyFrom samples a uniform polynomial using the given randomness source.
// Uses rejection sampling to avoid modular bias.
//
// For efficiency, reads entropy in bulk and uses the largest rejection bound
// that fits in the sample width. For q=257 with 16-bit samples: accept if
// v < floor(65536/257)*257 = 65535, then v % 257. Acceptance rate ~99.998%.
func UniformPolyFrom(p Params, rng io.Reader) *Poly {
	poly := New(p)
	q := p.Q

	if q <= (1 << 16) {
		// 16-bit sampling path (q fits in uint16).
		rejBound := uint32((65536 / uint32(q)) * uint32(q)) // largest multiple of q ≤ 65536

		// Batch read: need ~N samples, each 2 bytes. Over-allocate slightly
		// for the rare rejection.
		batchSize := p.N + 8 // small margin for rejections
		buf := make([]byte, batchSize*2)
		if _, err := io.ReadFull(rng, buf); err != nil {
			panic("ring: randomness source failed: " + err.Error())
		}
		pos := 0
		for i := range poly.Coeffs {
			for {
				if pos+2 > len(buf) {
					// Ran out of batch (very rare). Refill.
					buf = make([]byte, 16)
					io.ReadFull(rng, buf)
					pos = 0
				}
				v := uint32(binary.LittleEndian.Uint16(buf[pos : pos+2]))
				pos += 2
				if v < rejBound {
					poly.Coeffs[i] = v % q
					break
				}
			}
		}
	} else {
		// 32-bit sampling path (q > 16 bits).
		rejBound := (uint64(1)<<32 / uint64(q)) * uint64(q)

		batchSize := p.N + 8
		buf := make([]byte, batchSize*4)
		if _, err := io.ReadFull(rng, buf); err != nil {
			panic("ring: randomness source failed: " + err.Error())
		}
		pos := 0
		for i := range poly.Coeffs {
			for {
				if pos+4 > len(buf) {
					buf = make([]byte, 32)
					io.ReadFull(rng, buf)
					pos = 0
				}
				v := uint64(binary.LittleEndian.Uint32(buf[pos : pos+4]))
				pos += 4
				if v < rejBound {
					poly.Coeffs[i] = uint32(v % uint64(q))
					break
				}
			}
		}
	}
	return poly
}

// TernaryPoly samples a polynomial with coefficients in {-1, 0, 1}.
// Each coefficient is 0 with probability 1/2, and ±1 with probability 1/4 each.
func TernaryPoly(p Params) *Poly {
	return TernaryPolyFrom(p, rand.Reader)
}

// TernaryPolyFrom samples a ternary polynomial from the given source.
func TernaryPolyFrom(p Params, rng io.Reader) *Poly {
	poly := New(p)
	buf := make([]byte, p.N)
	if _, err := io.ReadFull(rng, buf); err != nil {
		panic("ring: randomness source failed: " + err.Error())
	}
	for i := range poly.Coeffs {
		b := buf[i]
		switch {
		case b < 64: // probability 1/4
			poly.Coeffs[i] = p.Q - 1 // -1 mod Q
		case b < 128: // probability 1/4
			poly.Coeffs[i] = 1
		default: // probability 1/2
			poly.Coeffs[i] = 0
		}
	}
	return poly
}

// SmallPoly samples a polynomial with coefficients in [-bound, bound].
// Uses centered representation: negative values stored as Q - |v|.
func SmallPoly(p Params, bound uint32) *Poly {
	return SmallPolyFrom(p, bound, rand.Reader)
}

// SmallPolyFrom samples a small polynomial from the given source.
func SmallPolyFrom(p Params, bound uint32, rng io.Reader) *Poly {
	poly := New(p)
	width := 2*bound + 1
	buf := make([]byte, 2)
	for i := range poly.Coeffs {
		for {
			if _, err := io.ReadFull(rng, buf); err != nil {
				panic("ring: randomness source failed: " + err.Error())
			}
			v := uint32(binary.LittleEndian.Uint16(buf))
			if v < (65535/width)*width { // rejection sampling for uniform
				v = v % width
				if v <= bound {
					poly.Coeffs[i] = v
				} else {
					poly.Coeffs[i] = p.Q - (width - v)
				}
				break
			}
		}
	}
	return poly
}

// CBDPoly samples from a centered binomial distribution with parameter eta.
// CBD_eta produces coefficients in [-eta, eta].
// This is the noise distribution used in Kyber/ML-KEM.
func CBDPoly(p Params, eta int) *Poly {
	return CBDPolyFrom(p, eta, rand.Reader)
}

// CBDPolyFrom samples CBD from the given source.
func CBDPolyFrom(p Params, eta int, rng io.Reader) *Poly {
	poly := New(p)
	bytesNeeded := (eta * 2 * p.N) / 8
	buf := make([]byte, bytesNeeded)
	if _, err := io.ReadFull(rng, buf); err != nil {
		panic("ring: randomness source failed: " + err.Error())
	}

	bitIdx := 0
	getBit := func() uint32 {
		bytePos := bitIdx / 8
		bitPos := uint(bitIdx % 8)
		bitIdx++
		return uint32((buf[bytePos] >> bitPos) & 1)
	}

	for i := range poly.Coeffs {
		var a, b uint32
		for j := range eta {
			_ = j
			a += getBit()
		}
		for j := range eta {
			_ = j
			b += getBit()
		}
		if a >= b {
			poly.Coeffs[i] = a - b
		} else {
			poly.Coeffs[i] = p.Q - (b - a)
		}
	}
	return poly
}

// Norm computes the infinity norm (max absolute coefficient) of a polynomial.
// Coefficients are interpreted as centered: values > Q/2 are negative.
func Norm(a *Poly) uint32 {
	q := a.params.Q
	half := q / 2
	maxVal := uint32(0)
	for _, c := range a.Coeffs {
		v := c
		if v > half {
			v = q - v // |c| for negative values
		}
		if v > maxVal {
			maxVal = v
		}
	}
	return maxVal
}

// Serialize packs polynomial coefficients into a byte slice.
// Each coefficient is stored as ceil(log2(Q)) bits, little-endian packed.
func Serialize(a *Poly) []byte {
	q := a.params.Q
	bits := 0
	for v := q - 1; v > 0; v >>= 1 {
		bits++
	}

	totalBits := bits * len(a.Coeffs)
	out := make([]byte, (totalBits+7)/8)

	bitPos := 0
	for _, c := range a.Coeffs {
		for b := 0; b < bits; b++ {
			if c&(1<<uint(b)) != 0 {
				out[bitPos/8] |= 1 << uint(bitPos%8)
			}
			bitPos++
		}
	}
	return out
}

// Deserialize unpacks a byte slice into a polynomial.
func Deserialize(p Params, data []byte) *Poly {
	q := p.Q
	bits := 0
	for v := q - 1; v > 0; v >>= 1 {
		bits++
	}

	poly := New(p)
	bitPos := 0
	for i := range poly.Coeffs {
		var v uint32
		for b := 0; b < bits; b++ {
			if bitPos/8 < len(data) && data[bitPos/8]&(1<<uint(bitPos%8)) != 0 {
				v |= 1 << uint(b)
			}
			bitPos++
		}
		if v >= q {
			v = q - 1 // clamp (shouldn't happen with well-formed data)
		}
		poly.Coeffs[i] = v
	}
	return poly
}