package services

import "sync"

// archiveTaskGate tracks work that may retain a *pvf.Archive after it has
// released core.mu. Closing the archive blocks new tasks, cancels the active
// ones, and waits for their references to disappear before releasing memory.
type archiveTaskGate struct {
	mu      sync.Mutex
	cond    *sync.Cond
	closing bool
	nextID  uint64
	tasks   map[uint64]chan struct{}
}

func newArchiveTaskGate() archiveTaskGate {
	// Initialize the condition lazily, after the gate has reached its final
	// address. Returning a cond bound to a local mutex copies the wrong lock.
	return archiveTaskGate{tasks: make(map[uint64]chan struct{})}
}

func (g *archiveTaskGate) ensureLocked() {
	if g.cond != nil {
		return
	}
	g.tasks = make(map[uint64]chan struct{})
	g.cond = sync.NewCond(&g.mu)
}

func (g *archiveTaskGate) begin() func() {
	g.mu.Lock()
	g.ensureLocked()
	for g.closing {
		g.cond.Wait()
	}
	return g.beginLocked()
}

// tryBegin lets an existing task register follow-up work without waiting for
// closeArchive, which may itself be waiting for that task to finish.
func (g *archiveTaskGate) tryBegin() (func(), bool) {
	g.mu.Lock()
	g.ensureLocked()
	if g.closing {
		g.mu.Unlock()
		return nil, false
	}
	return g.beginLocked(), true
}

func (g *archiveTaskGate) beginLocked() func() {
	g.nextID++
	id := g.nextID
	done := make(chan struct{})
	g.tasks[id] = done
	g.mu.Unlock()

	var once sync.Once
	return func() {
		once.Do(func() {
			g.mu.Lock()
			if current, ok := g.tasks[id]; ok {
				delete(g.tasks, id)
				close(current)
			}
			g.mu.Unlock()
		})
	}
}

func (g *archiveTaskGate) beginClose() []<-chan struct{} {
	g.mu.Lock()
	g.ensureLocked()
	for g.closing {
		g.cond.Wait()
	}
	g.closing = true
	waitFor := make([]<-chan struct{}, 0, len(g.tasks))
	for _, done := range g.tasks {
		waitFor = append(waitFor, done)
	}
	g.mu.Unlock()
	return waitFor
}

func (g *archiveTaskGate) endClose() {
	g.mu.Lock()
	g.ensureLocked()
	g.closing = false
	g.cond.Broadcast()
	g.mu.Unlock()
}