package crypto

import (
	"encoding/binary"
	"errors"
)

// Shadow compression wire format magic bytes.
var shadowMagic = [4]byte{'S', 'D', 'W', '1'}

// MarshalShadowCompressed encodes a ShadowCompressed into the binary
// wire format.
//
// Layout:
//
//	[4B]   Magic "SDW1"
//	[4B]   OrigLen (uint32 LE)
//	[4B]   TokenCount (uint32 LE)
//	[2B]   EpochDecExp (uint16 LE)
//	[2B]   EpochBinExp (uint16 LE)
//	[4B]   PrimaryLen (uint32 LE)
//	[4B]   ShadowLen (uint32 LE)
//	[4B]   ShadowPayloadLen (uint32 LE)
//	[256B] FreqTable (64 × uint32 LE)
//	[56B]  ContentHash (Hamadryad)
//	[PrimaryLen B] Huffman bitstream
//	[ShadowLen B]  Shadow permutation indices
func MarshalShadowCompressed(sc *ShadowCompressed) ([]byte, error) {
	if sc == nil {
		return nil, errors.New("crypto: nil shadow compressed")
	}

	headerSize := 4 + 4 + 4 + 2 + 2 + 4 + 4 + 4 + 256 + HamBytes // 340
	totalSize := headerSize + len(sc.Primary) + len(sc.Shadow)
	buf := make([]byte, totalSize)
	pos := 0

	// Magic.
	copy(buf[pos:], shadowMagic[:])
	pos += 4

	// OrigLen.
	binary.LittleEndian.PutUint32(buf[pos:], uint32(sc.OrigLen))
	pos += 4

	// TokenCount.
	binary.LittleEndian.PutUint32(buf[pos:], uint32(sc.TokenCount))
	pos += 4

	// EpochDecExp.
	binary.LittleEndian.PutUint16(buf[pos:], uint16(sc.EpochDec))
	pos += 2

	// EpochBinExp.
	binary.LittleEndian.PutUint16(buf[pos:], uint16(sc.EpochBin))
	pos += 2

	// PrimaryLen.
	binary.LittleEndian.PutUint32(buf[pos:], uint32(len(sc.Primary)))
	pos += 4

	// ShadowLen.
	binary.LittleEndian.PutUint32(buf[pos:], uint32(len(sc.Shadow)))
	pos += 4

	// ShadowPayloadLen.
	binary.LittleEndian.PutUint32(buf[pos:], uint32(sc.ShadowPayloadLen))
	pos += 4

	// FreqTable: 64 × uint32.
	for i := range 64 {
		binary.LittleEndian.PutUint32(buf[pos:], sc.FreqTable[i])
		pos += 4
	}

	// ContentHash.
	copy(buf[pos:], sc.ContentHash[:])
	pos += HamBytes

	// Primary stream.
	copy(buf[pos:], sc.Primary)
	pos += len(sc.Primary)

	// Shadow stream.
	copy(buf[pos:], sc.Shadow)

	return buf, nil
}

// UnmarshalShadowCompressed decodes a ShadowCompressed from the binary
// wire format.
func UnmarshalShadowCompressed(data []byte) (*ShadowCompressed, error) {
	headerSize := 4 + 4 + 4 + 2 + 2 + 4 + 4 + 4 + 256 + HamBytes
	if len(data) < headerSize {
		return nil, errors.New("crypto: shadow compressed data too short")
	}
	pos := 0

	// Magic.
	if data[0] != 'S' || data[1] != 'D' || data[2] != 'W' || data[3] != '1' {
		return nil, errors.New("crypto: invalid shadow compressed magic")
	}
	pos += 4

	sc := &ShadowCompressed{}

	// OrigLen.
	sc.OrigLen = int(binary.LittleEndian.Uint32(data[pos:]))
	pos += 4

	// TokenCount.
	sc.TokenCount = int(binary.LittleEndian.Uint32(data[pos:]))
	pos += 4

	// EpochDecExp.
	sc.EpochDec = int(binary.LittleEndian.Uint16(data[pos:]))
	pos += 2

	// EpochBinExp.
	sc.EpochBin = int(binary.LittleEndian.Uint16(data[pos:]))
	pos += 2

	// PrimaryLen.
	primaryLen := int(binary.LittleEndian.Uint32(data[pos:]))
	pos += 4

	// ShadowLen.
	shadowLen := int(binary.LittleEndian.Uint32(data[pos:]))
	pos += 4

	// ShadowPayloadLen.
	sc.ShadowPayloadLen = int(binary.LittleEndian.Uint32(data[pos:]))
	pos += 4

	// FreqTable.
	for i := range 64 {
		sc.FreqTable[i] = binary.LittleEndian.Uint32(data[pos:])
		pos += 4
	}

	// ContentHash.
	copy(sc.ContentHash[:], data[pos:pos+HamBytes])
	pos += HamBytes

	// Bounds check.
	if pos+primaryLen+shadowLen > len(data) {
		return nil, errors.New("crypto: shadow compressed data truncated")
	}

	// Primary stream.
	sc.Primary = make([]byte, primaryLen)
	copy(sc.Primary, data[pos:pos+primaryLen])
	pos += primaryLen

	// Shadow stream.
	sc.Shadow = make([]byte, shadowLen)
	copy(sc.Shadow, data[pos:pos+shadowLen])

	return sc, nil
}