package main

import (
	"container/heap"
	"slices"
)

const (
	// buildMax is the number of model builds that can run at once.
	buildMax = 8

	// buildLeafWindow is the number of unbuilt leaves that can be before a
	// leaf whose build starts.
	buildLeafWindow = 8
)

func newBuildQueue() *BuildQueue {
	return &BuildQueue{
		queued:  make(map[NodeKey]bool),
		running: make(map[NodeKey]bool),
		failed:  make(map[NodeKey]bool),
	}
}

// BuildQueue holds the nodes that wait for a model build. A leaf joins it when
// its message is too long to be its own node. A merge joins it when both of its
// children are built and their texts do not fit in one node. So the core never
// scans the tree for work.
//
// The builds start in the order of the messages: the node whose context ends
// first starts first. The context of leaf i ends before message i, and the
// context of a merge ends after its last message. At a tie the merge starts
// first. So a backlog of leaves does not stop the merges that the view needs to
// shrink. A leaf can start only when fewer than buildLeafWindow leaves before
// it are unbuilt. A failed build waits until the next message.
type BuildQueue struct {
	leaves  []int
	merges  BuildMergeHeap
	queued  map[NodeKey]bool
	running map[NodeKey]bool
	failed  map[NodeKey]bool
	retry   []NodeKey
}

// AddLeaf adds the leaf of message id. The IDs must come in increasing order.
func (q *BuildQueue) AddLeaf(id int) {
	q.leaves = append(q.leaves, id)
}

// AddMerge adds a merge whose children are built. A merge that is in the queue
// already stays as it is.
func (q *BuildQueue) AddMerge(key NodeKey) {
	if q.queued[key] {
		return
	}

	q.queued[key] = true
	heap.Push(&q.merges, key)
}

// Next returns the next node that can start, and marks it as running.
func (q *BuildQueue) Next() (NodeKey, bool) {
	if len(q.running) >= buildMax {
		return NodeKey{}, false
	}

	leaf, found := q.nextLeaf()
	if len(q.merges) > 0 && (!found || nodeEnd(q.merges[0].level, q.merges[0].index) <= leaf.index) {
		key := heap.Pop(&q.merges).(NodeKey)
		q.running[key] = true

		return key, true
	}

	if found {
		q.running[leaf] = true
	}

	return leaf, found
}

// Done removes a built node from the queue.
func (q *BuildQueue) Done(key NodeKey) {
	delete(q.running, key)
	if key.level > 0 {
		delete(q.queued, key)
		return
	}

	position := slices.Index(q.leaves, key.index)
	if position >= 0 {
		q.leaves = slices.Delete(q.leaves, position, position+1)
	}
}

// Fail keeps a node in the queue, but it cannot start until Retry.
func (q *BuildQueue) Fail(key NodeKey) {
	delete(q.running, key)
	if key.level > 0 {
		q.retry = append(q.retry, key)
		return
	}

	q.failed[key] = true
}

// Retry lets the failed nodes start again.
func (q *BuildQueue) Retry() {
	clear(q.failed)
	for _, key := range q.retry {
		heap.Push(&q.merges, key)
	}

	q.retry = q.retry[:0]
}

// Running returns the number of builds that run.
func (q *BuildQueue) Running() int {
	return len(q.running)
}

// nextLeaf returns the oldest leaf in the window that can start.
func (q *BuildQueue) nextLeaf() (NodeKey, bool) {
	for _, id := range q.leaves[:min(len(q.leaves), buildLeafWindow)] {
		key := NodeKey{0, id}
		if !q.running[key] && !q.failed[key] {
			return key, true
		}
	}

	return NodeKey{}, false
}

// BuildMergeHeap holds the merges that can start, by the end of their context.
type BuildMergeHeap []NodeKey

func (h BuildMergeHeap) Len() int {
	return len(h)
}

func (h BuildMergeHeap) Less(a, b int) bool {
	return nodeEnd(h[a].level, h[a].index) < nodeEnd(h[b].level, h[b].index)
}

func (h BuildMergeHeap) Swap(a, b int) {
	h[a], h[b] = h[b], h[a]
}

func (h *BuildMergeHeap) Push(value any) {
	*h = append(*h, value.(NodeKey))
}

func (h *BuildMergeHeap) Pop() any {
	old := *h
	key := old[len(old)-1]
	*h = old[:len(old)-1]

	return key
}