package crypto

import (
	"container/heap"
	"errors"

	"git.mleku.dev/mleku/dendrite/pkg/epoch"
	"git.mleku.dev/mleku/dendrite/pkg/ratio"
	"git.mleku.dev/mleku/dendrite/pkg/state"
)

// ShadowCompress compresses data using Huffman coding over hexagram
// tokens, with auxiliary data encoded in the S_3×S_3 permutation
// pair shadow channel.
//
// The primary stream is the Huffman-coded hexagram bitstream.
// The shadow stream encodes which of 36 possible permutation pairs
// was applied at each token position. The "home" pair (index 0) is
// determined by the epoch phase; deviations encode auxiliary bits.
//
// Shadow capacity: ~5 bits per hexagram token (log2(35) ≈ 5.1).
//
// Algorithm:
//  1. Convert data to hexagram tokens
//  2. Build frequency table and Huffman tree (exact rational weights)
//  3. For each token, select permutation pair based on auxiliary data
//  4. Huffman-encode the permuted tokens
//  5. Pack shadow stream (pair indices, 6 bits each)
func ShadowCompress(
	data []byte,
	auxiliary []byte,
	ep epoch.Epoch,
) (*ShadowCompressed, error) {
	if len(data) == 0 {
		return &ShadowCompressed{
			OrigLen:     0,
			EpochDec:    ep.DecExp,
			EpochBin:    ep.BinExp,
			ContentHash: Hash(nil),
		}, nil
	}

	// 1. Hexagram encoding.
	tokens := state.EncodeBytes(data)
	tokenCount := len(tokens)

	// 2. For each token, select permutation pair and apply shadow
	// decomposition. This must happen before building the frequency
	// table because Huffman coding operates on the permuted tokens.
	auxBits := bytesToBits(auxiliary)
	auxPos := 0

	shadowIndices := make([]uint8, tokenCount)
	permutedTokens := make([]state.Hexagram, tokenCount)

	for i, tok := range tokens {
		homeBin, homeDec := PhasePerms(ep, i)

		// Determine how many auxiliary bits to encode at this position.
		// We can encode up to 5 bits (values 1-35 map to 0-34 → 5 bits).
		var pairIdx uint8
		if auxPos < len(auxBits) {
			// Extract up to 5 bits of auxiliary data.
			val := uint8(0)
			bits := 0
			for bits < 5 && auxPos < len(auxBits) {
				val |= auxBits[auxPos] << uint(bits)
				auxPos++
				bits++
			}
			// Pair index 0 = home pair (no auxiliary data).
			// Pair indices 1-35 encode auxiliary values 0-34.
			if val < 35 {
				pairIdx = val + 1
			} else {
				pairIdx = 0 // overflow: use home pair
			}
		}

		shadowIndices[i] = pairIdx

		// Apply the selected permutation pair.
		binP, decP := PairFromIndex(pairIdx, homeBin, homeDec)
		permutedTokens[i] = ShadowDecompose(tok, binP, decP)
	}

	// 3. Build frequency table from permuted tokens.
	var freq [64]uint32
	for _, tok := range permutedTokens {
		freq[tok&0x3F]++
	}

	// 4. Build Huffman tree with exact rational weights.
	total := int64(tokenCount)
	tree := buildHuffmanTree(freq, total)
	codes := make(map[state.Hexagram]huffCode)
	buildCodes(tree, nil, codes)

	// 5. Huffman-encode the permuted tokens.
	primary := huffmanEncode(permutedTokens, codes)

	// 6. Pack shadow indices (6 bits each).
	shadow := packShadowIndices(shadowIndices)

	return &ShadowCompressed{
		Primary:          primary,
		Shadow:           shadow,
		FreqTable:        freq,
		EpochDec:         ep.DecExp,
		EpochBin:         ep.BinExp,
		OrigLen:          len(data),
		TokenCount:       tokenCount,
		ShadowPayloadLen: len(auxiliary),
		ContentHash:      Hash(data),
	}, nil
}

// ShadowDecompress recovers the original data and any embedded
// auxiliary payload from a shadow-compressed stream.
func ShadowDecompress(
	sc *ShadowCompressed,
	ep epoch.Epoch,
) (data []byte, auxiliary []byte, err error) {
	if sc == nil {
		return nil, nil, errors.New("crypto: nil shadow compressed")
	}
	if sc.OrigLen == 0 {
		return nil, nil, nil
	}

	// 1. Rebuild Huffman tree from frequency table.
	total := int64(sc.TokenCount)
	tree := buildHuffmanTree(sc.FreqTable, total)

	// 2. Huffman-decode the primary stream to get permuted tokens.
	permutedTokens := huffmanDecode(sc.Primary, tree, sc.TokenCount)

	// 3. Unpack shadow indices.
	shadowIndices := unpackShadowIndices(sc.Shadow, sc.TokenCount)

	// 4. Recover original tokens and extract auxiliary bits.
	tokens := make([]state.Hexagram, sc.TokenCount)
	var auxBits []uint8
	totalAuxBits := sc.ShadowPayloadLen * 8

	for i, ptok := range permutedTokens {
		homeBin, homeDec := PhasePerms(ep, i)
		pairIdx := shadowIndices[i]

		// Recover the permutation pair that was used.
		binP, decP := PairFromIndex(pairIdx, homeBin, homeDec)

		// Apply inverse permutation to recover original token.
		tokens[i] = ShadowRecompose(ptok, binP, decP)

		// Extract auxiliary bits from the pair index.
		// Only extract as many bits as the compress side packed:
		// the final group may have fewer than 5 bits.
		if pairIdx > 0 {
			val := pairIdx - 1
			for b := range 5 {
				if len(auxBits) >= totalAuxBits {
					break
				}
				auxBits = append(auxBits, (val>>uint(b))&1)
			}
		}
	}

	// 5. Decode hexagrams to bytes.
	data = state.DecodeHexagrams(tokens, sc.OrigLen)

	// 6. Convert auxiliary bits back to bytes.
	if sc.ShadowPayloadLen > 0 {
		auxiliary = bitsToBytes(auxBits, sc.ShadowPayloadLen)
	}

	// 7. Verify content hash.
	if Hash(data) != sc.ContentHash {
		return nil, nil, errors.New("crypto: shadow decompression integrity check failed")
	}

	return data, auxiliary, nil
}

// huffCode stores the Huffman code for a symbol.
type huffCode struct {
	bits   []uint8 // sequence of 0/1 bits, MSB first
	length int
}

// huffNode is a node in the Huffman tree with exact rational weights.
type huffNode struct {
	symbol state.Hexagram
	weight ratio.Ratio
	left   *huffNode
	right  *huffNode
	leaf   bool
}

// huffHeap implements heap.Interface for huffNode priority queue.
type huffHeap []*huffNode

func (h huffHeap) Len() int           { return len(h) }
func (h huffHeap) Less(i, j int) bool { return h[i].weight.Less(h[j].weight) }
func (h huffHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *huffHeap) Push(x any)        { *h = append(*h, x.(*huffNode)) }
func (h *huffHeap) Pop() any {
	old := *h
	n := len(old)
	item := old[n-1]
	*h = old[:n-1]
	return item
}

// buildHuffmanTree constructs a Huffman tree from frequency counts
// using exact rational arithmetic for all weight comparisons.
func buildHuffmanTree(freq [64]uint32, total int64) *huffNode {
	if total == 0 {
		total = 1
	}

	h := &huffHeap{}
	heap.Init(h)

	for i := range 64 {
		if freq[i] > 0 {
			heap.Push(h, &huffNode{
				symbol: state.Hexagram(i),
				weight: ratio.New(int64(freq[i]), total),
				leaf:   true,
			})
		}
	}

	// Handle edge cases: 0 or 1 symbols.
	if h.Len() == 0 {
		return &huffNode{leaf: true, weight: ratio.Zero}
	}
	if h.Len() == 1 {
		node := heap.Pop(h).(*huffNode)
		return &huffNode{
			weight: node.weight,
			left:   node,
			leaf:   false,
		}
	}

	// Standard Huffman construction.
	for h.Len() > 1 {
		a := heap.Pop(h).(*huffNode)
		b := heap.Pop(h).(*huffNode)
		heap.Push(h, &huffNode{
			weight: a.weight.Add(b.weight),
			left:   a,
			right:  b,
			leaf:   false,
		})
	}

	return heap.Pop(h).(*huffNode)
}

// buildCodes recursively builds the codebook from the Huffman tree.
func buildCodes(node *huffNode, prefix []uint8, codes map[state.Hexagram]huffCode) {
	if node == nil {
		return
	}
	if node.leaf {
		code := make([]uint8, len(prefix))
		copy(code, prefix)
		if len(code) == 0 {
			code = []uint8{0} // single-symbol edge case
		}
		codes[node.symbol] = huffCode{bits: code, length: len(code)}
		return
	}
	buildCodes(node.left, append(prefix, 0), codes)
	buildCodes(node.right, append(prefix, 1), codes)
}

// huffmanEncode encodes a sequence of hexagram tokens using Huffman codes.
func huffmanEncode(tokens []state.Hexagram, codes map[state.Hexagram]huffCode) []byte {
	var bits []uint8
	for _, tok := range tokens {
		code, ok := codes[tok]
		if !ok {
			// Token not in codebook — this can happen if the permuted
			// token wasn't in the original frequency table. Use the
			// identity code (longest possible).
			for range 6 {
				bits = append(bits, 0)
			}
			continue
		}
		bits = append(bits, code.bits...)
	}
	return bitsToBytesPacked(bits)
}

// huffmanDecode decodes a Huffman-coded bitstream back to hexagram tokens.
func huffmanDecode(data []byte, tree *huffNode, count int) []state.Hexagram {
	if tree == nil || count == 0 {
		return nil
	}

	bits := byteToBitsPacked(data)
	tokens := make([]state.Hexagram, 0, count)
	node := tree
	for _, bit := range bits {
		if len(tokens) >= count {
			break
		}
		if bit == 0 {
			node = node.left
		} else {
			node = node.right
		}
		if node == nil {
			break
		}
		if node.leaf {
			tokens = append(tokens, node.symbol)
			node = tree
		}
	}
	return tokens
}

// packShadowIndices packs 6-bit indices into a byte stream.
// Each index is 0-35, fitting in 6 bits.
func packShadowIndices(indices []uint8) []byte {
	totalBits := len(indices) * 6
	out := make([]byte, (totalBits+7)/8)
	bitPos := 0
	for _, idx := range indices {
		for b := range 6 {
			if idx&(1<<uint(b)) != 0 {
				out[bitPos/8] |= 1 << uint(bitPos%8)
			}
			bitPos++
		}
	}
	return out
}

// unpackShadowIndices unpacks 6-bit indices from a byte stream.
func unpackShadowIndices(data []byte, count int) []uint8 {
	indices := make([]uint8, count)
	bitPos := 0
	for i := range count {
		var val uint8
		for b := range 6 {
			byteIdx := bitPos / 8
			bitIdx := uint(bitPos % 8)
			if byteIdx < len(data) && data[byteIdx]&(1<<bitIdx) != 0 {
				val |= 1 << uint(b)
			}
			bitPos++
		}
		indices[i] = val
	}
	return indices
}

// bytesToBits converts a byte slice to individual bits (LSB first per byte).
func bytesToBits(data []byte) []uint8 {
	bits := make([]uint8, len(data)*8)
	for i, b := range data {
		for j := range 8 {
			bits[i*8+j] = (b >> uint(j)) & 1
		}
	}
	return bits
}

// bitsToBytes converts individual bits back to bytes (LSB first per byte).
func bitsToBytes(bits []uint8, length int) []byte {
	out := make([]byte, length)
	for i := range length {
		var b byte
		for j := range 8 {
			idx := i*8 + j
			if idx < len(bits) && bits[idx] != 0 {
				b |= 1 << uint(j)
			}
		}
		out[i] = b
	}
	return out
}

// bitsToBytesPacked packs a bit sequence into bytes (MSB first).
func bitsToBytesPacked(bits []uint8) []byte {
	out := make([]byte, (len(bits)+7)/8)
	for i, bit := range bits {
		if bit != 0 {
			out[i/8] |= 1 << uint(7-i%8)
		}
	}
	return out
}

// byteToBitsPacked unpacks bytes into individual bits (MSB first).
func byteToBitsPacked(data []byte) []uint8 {
	bits := make([]uint8, len(data)*8)
	for i, b := range data {
		for j := range 8 {
			bits[i*8+j] = (b >> uint(7-j)) & 1
		}
	}
	return bits
}