package instancectl

import (
	"context"
	"errors"
	"fmt"
	"net"
	"time"
)

// Legacy run start attempts (#292, Astra round 2 P1, the owner's ruling).
//
// The agent proves an aborted start's unit stopped through /stop before it
// gives the lease back. Cancelling its /start does not join this side's
// handler, which could still submit the systemd start after that /stop had
// answered, if it had not yet seen the disconnect; systemd's replacement of
// a queued start by a stop only covers jobs that already exist. So, for the
// original run unit (carlos-instance@, generation 0) only, start and stop
// are ordered per host here:
//
//   - every /start registers an attempt with a context /stop can cancel,
//     and checks that context immediately before it submits the start;
//   - a /stop fences the host: new starts are refused while it runs, every
//     attempt in flight is cancelled and joined (bounded), and only then is
//     the unit stopped, which replaces any job a joined attempt submitted;
//   - a start whose request was accepted at or before the moment a stop
//     fences the host (the accepted-connection watermark, connSeq) is an
//     earlier request, and is refused, however late its handler registers;
//     one accepted after it runs normally. Serve answers one request per
//     connection, so a connection's accept order is its request's.
//
// Once /stop has answered, no start from an earlier request is in flight or
// can be submitted afterwards. Generation units, seeds, sidecars and the
// claim verbs are untouched.

// legacyJoinBudget bounds how long a /stop waits for the start attempts it
// cancelled to return. One that is still running then may yet submit, so
// the stop answers an error and the agent keeps the lease.
const legacyJoinBudget = 10 * time.Second

// legacyAttempt is one in-flight legacy run start.
type legacyAttempt struct {
	cancel context.CancelFunc
	done   chan struct{}
}

// connSeqKey carries a connection's accept order in its request contexts.
type connSeqKey struct{}

// nextConnSeq hands out the next accept-order number (from 1).
func (s *Server) nextConnSeq() uint64 { return s.connSeq.Add(1) }

// connContext is Serve's ConnContext: the claim verbs' connection value,
// and the connection's accept order. ConnContext runs on the accept loop,
// once per connection, in accept order.
func (s *Server) connContext(ctx context.Context, c net.Conn) context.Context {
	return context.WithValue(requestConnContext(ctx, c), connSeqKey{}, s.nextConnSeq())
}

// requestSeq is a request's accept order; 0 when its connection was not
// stamped (a handler called outside Serve), which no fence refuses.
func requestSeq(ctx context.Context) uint64 {
	n, _ := ctx.Value(connSeqKey{}).(uint64)
	return n
}

// errLegacyFenced refuses a start while a stop of its host runs, or one
// whose request came before a stop of its host.
var errLegacyFenced = errors.New("a stop of this host's run unit has run since this start was asked for")

// beginLegacyStart registers a legacy run start of host for a request of
// accept order seq. Its context is parent's, cancellable by a /stop of the
// host; finish must be called when the handler returns.
func (s *Server) beginLegacyStart(parent context.Context, host string, seq uint64) (context.Context, func(), error) {
	s.legacyMu.Lock()
	defer s.legacyMu.Unlock()
	if s.legacyStopping[host] > 0 || (seq != 0 && seq <= s.legacyBarrier[host]) {
		return nil, nil, errLegacyFenced
	}
	ctx, cancel := context.WithCancel(parent)
	a := &legacyAttempt{cancel: cancel, done: make(chan struct{})}
	if s.legacyStarts == nil {
		s.legacyStarts = make(map[string]map[*legacyAttempt]struct{})
	}
	if s.legacyStarts[host] == nil {
		s.legacyStarts[host] = make(map[*legacyAttempt]struct{})
	}
	s.legacyStarts[host][a] = struct{}{}
	finish := func() {
		s.legacyMu.Lock()
		delete(s.legacyStarts[host], a)
		if len(s.legacyStarts[host]) == 0 {
			delete(s.legacyStarts, host)
		}
		s.legacyMu.Unlock()
		cancel()
		close(a.done)
	}
	return ctx, finish, nil
}

// fenceLegacyStarts is a /stop's first step for host: refuse new starts
// while the stop runs, refuse every start accepted at or before this moment
// (the accepted-connection watermark, read under legacyMu, so no start can
// register between the watermark and the snapshot), cancel the attempts in
// flight and join them, bounded by the join budget and ctx. The fence is
// per host and in memory only: release lifts the refusal of new starts when
// the stop is done, and only connections accepted before the watermark
// stay refused. An error means an attempt did not return in time and may
// still submit.
func (s *Server) fenceLegacyStarts(ctx context.Context, host string) (func(), error) {
	s.legacyMu.Lock()
	seq := s.connSeq.Load()
	if s.legacyStopping == nil {
		s.legacyStopping = make(map[string]int)
	}
	if s.legacyBarrier == nil {
		s.legacyBarrier = make(map[string]uint64)
	}
	s.legacyStopping[host]++
	if seq > s.legacyBarrier[host] {
		s.legacyBarrier[host] = seq
	}
	var inFlight []*legacyAttempt
	for a := range s.legacyStarts[host] {
		inFlight = append(inFlight, a)
	}
	s.legacyMu.Unlock()
	release := func() {
		s.legacyMu.Lock()
		if s.legacyStopping[host]--; s.legacyStopping[host] <= 0 {
			delete(s.legacyStopping, host)
		}
		s.legacyMu.Unlock()
	}
	budget := legacyJoinBudget
	if s.legacyJoin > 0 {
		budget = s.legacyJoin
	}
	timer := time.NewTimer(budget)
	defer timer.Stop()
	for _, a := range inFlight {
		a.cancel()
		select {
		case <-a.done:
		case <-timer.C:
			return release, fmt.Errorf("a start of %s is still in flight after %s", host, budget)
		case <-ctx.Done():
			return release, fmt.Errorf("a start of %s is still in flight: %w", host, ctx.Err())
		}
	}
	return release, nil
}