// Package chunk provides a heap of chunks so that larger in-order chunks can be
// processed more easily, e.g., for verifying ranged inclusion proofs and
// batching file writes.  For how to implement the heap interface, see:
// https://pkg.go.dev/container/heap#Interface
package chunk

import (
	"container/heap"
	"crypto/sha256"
)

// Chunk holds the SANs of a consecutive range of downloaded leaves.  The start
// index and associated leaf hashes are included so that they can be verified.
type Chunk struct {
	Start      int64               // index of first leaf in this chunk
	LeafHashes [][sha256.Size]byte // in-order leaf hashes in this chunk
	SANs       []string            // sans of all leaves in this chunk
}

// ChunkHeap is a min-heap of chunks wrt. to start indices.  Use TPush() and
// TPop() to to push/pop chunks without having to worry about type casts.
type ChunkHeap []*Chunk

// TPop pops a chunk from the heap
func (h *ChunkHeap) TPop() *Chunk {
	x := heap.Pop(h)
	return x.(*Chunk)
}

// TPush pushes a chunk to the heap
func (h *ChunkHeap) TPush(c *Chunk) {
	heap.Push(h, c)
}

// Sequence checks if the next pop would provide a chunk starting at index
// start.  If so, the heap is also re-organized so that the popped chunk will be
// as large as possible (i.e., in-order chunks are merged to a single chunk).
func (h *ChunkHeap) Sequence(start int64) bool {
	if len(*h) == 0 {
		return false
	}

	// See if we have the next in-order chunk ready
	s := h.TPop()
	if start != s.Start {
		h.TPush(s)
		return false
	}

	// See if we can make a larger in-order chunk
	for len(*h) > 0 {
		c := h.TPop()
		if c.Start != s.Start+int64(len(s.LeafHashes)) {
			h.TPush(c)
			break
		}

		s.LeafHashes = append(s.LeafHashes, c.LeafHashes...)
		s.SANs = append(s.SANs, c.SANs...)
	}

	// Put back the largest in-order chunk we have
	h.TPush(s)
	return true
}

//
// Implement the heap interface
//

func (h ChunkHeap) Len() int           { return len(h) }
func (h ChunkHeap) Less(i, j int) bool { return h[i].Start < h[j].Start }
func (h ChunkHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *ChunkHeap) Push(x any) {
	*h = append(*h, x.(*Chunk))
}

func (h *ChunkHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	old[n-1] = nil // avoid memory leak
	*h = old[:n-1]
	return x
}