package dispatcher

import (
	"context"
	"errors"
	"sync"
)

var errProcessExitCallbackAdmissionClosed = errors.New("dispatcher process exit callback admission is closed")

// processExitCallbackOwner exists for the entire lifetime of a real Process,
// not merely while its callback is running. That makes shutdown independent of
// session/execution index membership: those indexes may be cleared at the
// beginning of onExit while durable callback work is still pending.
type processExitCallbackOwner struct {
	done     chan struct{}
	doneOnce sync.Once
}

func (d *Dispatcher) beginProcessExitCallbackOwner() (*processExitCallbackOwner, error) {
	if d == nil {
		return nil, errProcessExitCallbackAdmissionClosed
	}
	d.processExitCallbackMu.Lock()
	defer d.processExitCallbackMu.Unlock()
	if d.processExitCallbackClosed {
		return nil, errProcessExitCallbackAdmissionClosed
	}
	if d.processExitCallbackOwners == nil {
		d.processExitCallbackOwners = make(map[*processExitCallbackOwner]struct{})
	}
	owner := &processExitCallbackOwner{done: make(chan struct{})}
	d.processExitCallbackOwners[owner] = struct{}{}
	return owner, nil
}

func (d *Dispatcher) finishProcessExitCallbackOwner(owner *processExitCallbackOwner) {
	if owner == nil {
		return
	}
	if d != nil {
		d.processExitCallbackMu.Lock()
		delete(d.processExitCallbackOwners, owner)
		d.processExitCallbackMu.Unlock()
	}
	owner.doneOnce.Do(func() { close(owner.done) })
}

func (d *Dispatcher) closeProcessExitCallbackAdmission() {
	if d == nil {
		return
	}
	d.processExitCallbackMu.Lock()
	d.processExitCallbackClosed = true
	d.processExitCallbackMu.Unlock()
}

func (d *Dispatcher) waitForProcessExitCallbackOwners(ctx context.Context) error {
	if d == nil {
		return nil
	}
	d.processExitCallbackMu.Lock()
	owners := make([]<-chan struct{}, 0, len(d.processExitCallbackOwners))
	for owner := range d.processExitCallbackOwners {
		owners = append(owners, owner.done)
	}
	d.processExitCallbackMu.Unlock()
	return waitForSessionStartOwners(ctx, owners)
}