package gnarlring

// CommitmentFrame is the wire format for a single child commitment (63 bytes).
//
//	Offset  Size  Field
//	0       1     Index (uint8, 0..26)
//	1       31    PubKey (Poly27, 9-bit unsigned)
//	32      31    W (commitment, Poly27, 9-bit unsigned)
const CommitmentFrameSize = 1 + PolyBytes + PolyBytes // 63

func MarshalCommitmentFrame(cc *ChildCommitment) []byte {
	buf := make([]byte, CommitmentFrameSize)
	buf[0] = cc.Index
	copy(buf[1:1+PolyBytes], cc.PubKey.MarshalBinary())
	copy(buf[1+PolyBytes:], cc.W.MarshalBinary())
	return buf
}

func UnmarshalCommitmentFrame(data []byte) *ChildCommitment {
	if len(data) < CommitmentFrameSize {
		return nil
	}
	cc := &ChildCommitment{Index: data[0]}
	var err error
	cc.PubKey, err = UnmarshalBinary(data[1 : 1+PolyBytes])
	if err != nil {
		return nil
	}
	cc.W, err = UnmarshalBinary(data[1+PolyBytes : CommitmentFrameSize])
	if err != nil {
		return nil
	}
	return cc
}

// EpochFrame is the full epoch broadcast (≈1783 bytes).
//
//	Offset  Size    Field
//	0       8       Counter (uint64 LE)
//	8       31      RootPK (Poly27, 9-bit unsigned)
//	39      63×27   Frames [N]CommitmentFrame
//	1740    47      RootSig (16B salt + 31B s2 at 9-bit signed)
const EpochFrameHeaderSize = 8 + PolyBytes                    // 39
const EpochFrameSize = EpochFrameHeaderSize + N*CommitmentFrameSize + sigBytes // 39 + 1701 + 47 = 1787

func MarshalEpochFrame(es *EpochState) []byte {
	buf := make([]byte, EpochFrameSize)
	putUint64LE(buf[0:8], es.Counter)
	copy(buf[8:8+PolyBytes], es.Coordinator.MarshalBinary())

	off := EpochFrameHeaderSize
	for i := 0; i < N; i++ {
		if es.Agg.Children[i] != nil {
			copy(buf[off:off+CommitmentFrameSize], MarshalCommitmentFrame(es.Agg.Children[i]))
		}
		off += CommitmentFrameSize
	}

	if es.RootSig != nil {
		copy(buf[off:], es.RootSig.MarshalBinary())
	}
	return buf
}

func UnmarshalEpochFrame(data []byte) *EpochState {
	if len(data) < EpochFrameSize {
		return nil
	}
	c := getUint64LE(data[0:8])
	pk, err := UnmarshalBinary(data[8 : 8+PolyBytes])
	if err != nil {
		return nil
	}

	es := StartEpoch(c, &NTRUPublicKey{H: pk})
	off := EpochFrameHeaderSize
	for i := 0; i < N; i++ {
		cc := UnmarshalCommitmentFrame(data[off : off+CommitmentFrameSize])
		if cc != nil {
			es.Agg.Add(cc)
		}
		off += CommitmentFrameSize
	}

	sig, err := UnmarshalNTRUSig(data[off : off+sigBytes])
	if err == nil {
		es.RootSig = sig
		es.finalized = true
	}
	return es
}

// EpochCheckFrame is the compact relay-path proof (113 bytes).
//
//	Offset  Size  Field
//	0       8     Counter (uint64 LE)
//	8       31    RootPK (Poly27, 9-bit unsigned)
//	39      27    WCompressed (hash of all w_i + epoch, 27 bytes)
//	66      47    RootSig (16B salt + 31B s2)
//
// Does NOT carry the full w_i list. The verifier must have cached the
// epoch frame or trust the WCompressed binding.
const EpochCheckFrameSize = 8 + PolyBytes + 27 + sigBytes // 113

func MarshalEpochCheckFrame(es *EpochState) []byte {
	buf := make([]byte, EpochCheckFrameSize)
	putUint64LE(buf[0:8], es.Counter)
	copy(buf[8:8+PolyBytes], es.Coordinator.MarshalBinary())
	copy(buf[8+PolyBytes:8+PolyBytes+27], es.Agg.WCompressed(es.Counter))
	if es.RootSig != nil {
		copy(buf[8+PolyBytes+27:], es.RootSig.MarshalBinary())
	}
	return buf
}

func UnmarshalEpochCheckFrame(data []byte) (*EpochCheckFrame, error) {
	if len(data) < EpochCheckFrameSize {
		return nil, errShortData
	}
	cf := &EpochCheckFrame{Counter: getUint64LE(data[0:8])}
	var err error
	cf.PK, err = UnmarshalBinary(data[8 : 8+PolyBytes])
	if err != nil {
		return nil, err
	}
	copy(cf.WCompressed[:], data[8+PolyBytes:8+PolyBytes+27])
	cf.Sig, err = UnmarshalNTRUSig(data[8+PolyBytes+27 : EpochCheckFrameSize])
	if err != nil {
		return nil, err
	}
	return cf, nil
}

// EpochCheckFrame is a deserialized compact frame.
type EpochCheckFrame struct {
	Counter     uint64
	PK          *Poly27
	WCompressed [27]byte
	Sig         *NTRUSignature
}

func putUint64LE(buf []byte, v uint64) {
	buf[0] = byte(v)
	buf[1] = byte(v >> 8)
	buf[2] = byte(v >> 16)
	buf[3] = byte(v >> 24)
	buf[4] = byte(v >> 32)
	buf[5] = byte(v >> 40)
	buf[6] = byte(v >> 48)
	buf[7] = byte(v >> 56)
}

func getUint64LE(buf []byte) uint64 {
	return uint64(buf[0]) | uint64(buf[1])<<8 | uint64(buf[2])<<16 |
		uint64(buf[3])<<24 | uint64(buf[4])<<32 | uint64(buf[5])<<40 |
		uint64(buf[6])<<48 | uint64(buf[7])<<56
}