package rollout

import (
	"context"
	"sync"
	"time"
)

// ackWriteBound bounds each config acknowledgement the armed scheduler
// writes (Task 11 rule 3, Astra R1 Flag triage T8): a store that hangs holds
// up the next ack this long, never the scheduler. Flag off, config sync's
// own writer (host.configAckFunc) is unchanged (plan review #11).
const ackWriteBound = 10 * time.Second

// ackJob is one acknowledgement for the writer: seq is its place in the
// order of the config transitions that released it, taken under Manager.mu
// in the same critical section (Manager.ackLocked), and write is the store
// write, handed the bounded context.
type ackJob struct {
	seq   uint64
	write func(context.Context)
}

// ackWriter is Manager.acks: the armed scheduler's acknowledgement writes,
// off the paths that release them. A config activation releases its held ack
// in done() before the attempt's slot frees, and in the scheduler loop for a
// sleeper (configSlept); neither may wait on the store, so they submit here
// and one drain goroutine writes, with no lock held while it writes (Task 11
// rule 3).
//
// The order is fixed at the transition, not at the write (plan review #10):
// a host's acks only ever replace each other, so the writer keeps the newest
// pending one per host, by sequence, and drops an older one, whether it
// arrives after a newer one is pending or after a newer one was written.
// An ack written after a newer one would name a revision that is no longer
// the delivered one (spec §8.1). Hosts are written in the order their
// pending ack was first queued. Submission never blocks.
type ackWriter struct {
	ctx  context.Context
	logf func(format string, args ...any)
	kick chan struct{}
	done chan struct{} // closed when the drain goroutine has returned

	mu      sync.Mutex
	pending map[string]ackJob
	order   []string          // hosts with a pending ack, first queued first
	newest  map[string]uint64 // per host, the newest sequence accepted
	closing bool              // Close: drain what is pending, then return
	stop    bool              // return after the write in progress
	writing bool
	// onSubmit, when set, runs at every accepted submission, with no lock
	// held: a test's window into the moment an ack is released.
	onSubmit func(host string)
}

// newAckWriter starts the drain goroutine. It runs until ctx ends or close
// returns; a write in progress when ctx ends is finished (its own context is
// ctx's, so the store call returns as soon as it notices) and nothing more is
// written.
func newAckWriter(ctx context.Context, logf func(format string, args ...any)) *ackWriter {
	w := &ackWriter{ctx: ctx, logf: logf, kick: make(chan struct{}, 1), done: make(chan struct{}),
		pending: map[string]ackJob{}, newest: map[string]uint64{}}
	go w.drain()
	return w
}

// submit queues j as host's ack unless an ack with a newer sequence was
// already queued or written for host. It never blocks.
func (w *ackWriter) submit(host string, j ackJob) {
	w.mu.Lock()
	if w.stop || j.seq <= w.newest[host] {
		w.mu.Unlock()
		return
	}
	w.newest[host] = j.seq
	if _, queued := w.pending[host]; !queued {
		w.order = append(w.order, host)
	}
	w.pending[host] = j
	hook := w.onSubmit
	w.mu.Unlock()
	if hook != nil {
		hook(host)
	}
	select {
	case w.kick <- struct{}{}:
	default:
	}
}

// next takes the first queued host's ack. ok is false when there is none,
// or the writer is to stop.
func (w *ackWriter) next() (host string, j ackJob, ok bool) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.stop || len(w.order) == 0 || w.ctx.Err() != nil {
		w.writing = false
		return "", ackJob{}, false
	}
	host, w.order = w.order[0], w.order[1:]
	j = w.pending[host]
	delete(w.pending, host)
	w.writing = true
	return host, j, true
}

func (w *ackWriter) drain() {
	defer close(w.done)
	for {
		for {
			_, j, ok := w.next()
			if !ok {
				break
			}
			ctx, cancel := context.WithTimeout(w.ctx, ackWriteBound)
			j.write(ctx)
			cancel()
		}
		w.mu.Lock()
		finished := w.stop || (w.closing && len(w.order) == 0)
		w.mu.Unlock()
		if finished {
			return
		}
		select {
		case <-w.ctx.Done():
			return
		case <-w.kick:
		}
	}
}

// idle reports whether nothing is pending or being written.
func (w *ackWriter) idle() bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	return len(w.order) == 0 && !w.writing
}

// close writes what is pending and stops the writer, waiting until ctx
// ends. Once ctx ends the writer stops after the write in progress, and
// close returns without waiting for it: what is still pending then is not
// written. An ack submitted after close is dropped.
func (w *ackWriter) close(ctx context.Context) {
	w.mu.Lock()
	w.closing = true
	w.mu.Unlock()
	select {
	case w.kick <- struct{}{}:
	default:
	}
	select {
	case <-w.done:
	case <-ctx.Done():
	}
	w.mu.Lock()
	w.stop = true
	left := len(w.order)
	w.mu.Unlock()
	if left > 0 {
		w.logf("rollout: %d config acknowledgement(s) not written before the agent stopped", left)
	}
}