package gnarlring

import (
	"crypto/rand"
	"io"
	"math"
)

// ChildCommitment holds one child's per-epoch short-vector commitment.
// The child produces w = H_i * z (z freshly sampled Gaussian), and the
// coordinator collects all 27 w_i to hash into the epoch target.
type ChildCommitment struct {
	Index  uint8   // child position 0..26
	PubKey *Poly27 // child's NTRU public key H_i (31 bytes serialized)
	W      *Poly27 // commitment w_i = H_i * z_i (31 bytes serialized)
}

// NewChildCommitment creates a commitment: samples fresh z, computes w = pk * z.
func NewChildCommitment(index int, pk *NTRUPublicKey, rng io.Reader) *ChildCommitment {
	if rng == nil {
		rng = rand.Reader
	}
	gs := NewGaussSamplerFrom(DefaultSigma(), rng)
	z := gs.SamplePoly()
	w := Mul(pk.H, z)
	return &ChildCommitment{
		Index:  uint8(index),
		PubKey: pk.H,
		W:      w,
	}
}

// AggregatedCommitment collects up to 27 child commitments by index.
type AggregatedCommitment struct {
	Children [N]*ChildCommitment // nil for unused slots
	Count    int
}

// NewAggregatedCommitment returns an empty aggregation.
func NewAggregatedCommitment() *AggregatedCommitment {
	return &AggregatedCommitment{}
}

// Add inserts a child commitment at its index. Returns false on conflict.
func (ac *AggregatedCommitment) Add(cc *ChildCommitment) bool {
	if cc.Index >= N {
		return false
	}
	if ac.Children[cc.Index] != nil {
		return false
	}
	ac.Children[cc.Index] = cc
	ac.Count++
	return true
}

// Remove clears a slot.
func (ac *AggregatedCommitment) Remove(index uint8) {
	if index < N && ac.Children[index] != nil {
		ac.Children[index] = nil
		ac.Count--
	}
}

// IsComplete reports whether all N slots are filled.
func (ac *AggregatedCommitment) IsComplete() bool {
	return ac.Count == N
}

// Target produces the sparse challenge polynomial for the coordinator to sign.
// Hashes epoch counter + all w_i (zero for empty slots) + message. Returns a
// dense Poly27 suitable as ffSampling target.
func (ac *AggregatedCommitment) Target(msg []byte, epoch uint64) *Poly27 {
	input := make([]byte, 0, 16+N*PolyBytes+len(msg))
	input = append(input, []byte("gnarl-epoch-v1")...)
	input = appendUint64LE(input, epoch)
	for i := 0; i < N; i++ {
		if ac.Children[i] != nil {
			input = append(input, ac.Children[i].W.MarshalBinary()...)
		} else {
			input = append(input, make([]byte, PolyBytes)...)
		}
	}
	input = append(input, msg...)

	return hashBytesToPoly(input)
}

// WCompressed returns the GMid-style hash of all w_i for compact verification.
func (ac *AggregatedCommitment) WCompressed(epoch uint64) []byte {
	input := make([]byte, 0, 16+N*PolyBytes)
	input = append(input, []byte("gnarl-wcomp-v1")...)
	input = appendUint64LE(input, epoch)
	for i := 0; i < N; i++ {
		if ac.Children[i] != nil {
			input = append(input, ac.Children[i].W.MarshalBinary()...)
		} else {
			input = append(input, make([]byte, PolyBytes)...)
		}
	}
	return hashBytes(input)
}

func appendUint64LE(buf []byte, v uint64) []byte {
	var tmp [8]byte
	tmp[0] = byte(v)
	tmp[1] = byte(v >> 8)
	tmp[2] = byte(v >> 16)
	tmp[3] = byte(v >> 24)
	tmp[4] = byte(v >> 32)
	tmp[5] = byte(v >> 40)
	tmp[6] = byte(v >> 48)
	tmp[7] = byte(v >> 56)
	return append(buf, tmp[:]...)
}

// DefaultSigma returns the standard Gaussian sigma for the gnarl ring.
func DefaultSigma() float64 {
	return math.Sqrt(float64(N)) * 2.0
}

// hashBytes returns a 27-byte hash of input using FNV-style state chaining.
func hashBytes(input []byte) []byte {
	var state uint64 = 14695981039346656037
	out := make([]byte, 27)
	for i := 0; i < 27; i++ {
		for _, b := range input {
			state ^= uint64(b)
			state *= 1099511628211
		}
		state ^= uint64(i)
		state *= 1099511628211
		out[i] = byte(state >> 32)
	}
	return out
}

// hashBytesToPoly returns a dense Poly27 derived from input.
func hashBytesToPoly(input []byte) *Poly27 {
	var state uint64 = 14695981039346656037
	c := NewPoly27()
	for i := 0; i < N; i++ {
		for _, b := range input {
			state ^= uint64(b)
			state *= 1099511628211
		}
		state ^= uint64(i)
		state *= 1099511628211
		c.Coeffs[i] = uint16(state % uint64(Q))
	}
	return c
}

// salt generates a random 16-byte salt.
func salt16() ([16]byte, error) {
	var s [16]byte
	_, err := rand.Read(s[:])
	return s, err
}