package wire

import (
	"context"
	"sync"
)

// connectionRun is the exact owner domain for one established Session Link.
//
// Every goroutine that can retain connection-scoped context, mutate Protocol
// state on behalf of that connection, or enqueue a frame into the shared
// sendCh must be admitted here before it starts. closeAdmission atomically
// rejects future work, cancels every admitted owner, and returns a stable set
// of done channels. That explicit registry avoids the WaitGroup Add/Wait race:
// once the snapshot is taken, no unobserved owner can appear behind it.
type connectionRun struct {
	mu     sync.Mutex
	ctx    context.Context
	cancel context.CancelFunc
	closed bool
	owners map[*connectionRunOwner]struct{}
}

type connectionRunOwner struct {
	cancel   context.CancelFunc
	done     chan struct{}
	doneOnce sync.Once
}

type connectionRunContextKey struct{}

func newConnectionRun(parent context.Context) *connectionRun {
	if parent == nil {
		parent = context.Background()
	}
	baseCtx, cancel := context.WithCancel(parent)
	run := &connectionRun{
		cancel: cancel,
		owners: make(map[*connectionRunOwner]struct{}),
	}
	// Descendants retain the exact run identity, which lets deeply nested wire
	// helpers register work without threading another lifecycle parameter through
	// every protocol handler interface.
	run.ctx = context.WithValue(baseCtx, connectionRunContextKey{}, run)
	return run
}

func (r *connectionRun) context() context.Context {
	if r == nil || r.ctx == nil {
		return context.Background()
	}
	return r.ctx
}

func connectionRunFromContext(ctx context.Context) *connectionRun {
	if ctx == nil {
		return nil
	}
	run, _ := ctx.Value(connectionRunContextKey{}).(*connectionRun)
	return run
}

// start admits one owner before launching its goroutine. taskCtx preserves the
// caller's narrower deadline while closeAdmission can still cancel it directly.
func (r *connectionRun) start(taskCtx context.Context, task func(context.Context)) bool {
	if r == nil || task == nil {
		return false
	}
	if taskCtx == nil {
		taskCtx = r.context()
	}

	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return false
	}
	ownerCtx, cancelOwner := context.WithCancel(taskCtx)
	owner := &connectionRunOwner{
		cancel: cancelOwner,
		done:   make(chan struct{}),
	}
	r.owners[owner] = struct{}{}
	r.mu.Unlock()

	go func() {
		defer r.retire(owner)
		task(ownerCtx)
	}()
	return true
}

func (r *connectionRun) retire(owner *connectionRunOwner) {
	if r == nil || owner == nil {
		return
	}
	owner.doneOnce.Do(func() {
		if owner.cancel != nil {
			owner.cancel()
		}
		r.mu.Lock()
		delete(r.owners, owner)
		r.mu.Unlock()
		close(owner.done)
	})
}

// closeAdmission is idempotent. It closes task publication and cancels the
// exact owner set while holding the same mutex used by start, then returns the
// stable join snapshot after dropping the mutex.
func (r *connectionRun) closeAdmission() []<-chan struct{} {
	if r == nil {
		return nil
	}
	r.mu.Lock()
	if !r.closed {
		r.closed = true
		if r.cancel != nil {
			r.cancel()
		}
	}
	owners := make([]<-chan struct{}, 0, len(r.owners))
	for owner := range r.owners {
		if owner == nil {
			continue
		}
		if owner.cancel != nil {
			owner.cancel()
		}
		owners = append(owners, owner.done)
	}
	r.mu.Unlock()
	return owners
}

func waitForConnectionRunOwners(ctx context.Context, owners []<-chan struct{}) error {
	if ctx == nil {
		ctx = context.Background()
	}
	for _, done := range owners {
		if done == nil {
			continue
		}
		select {
		case <-done:
		case <-ctx.Done():
			return ctx.Err()
		}
	}
	return nil
}

// finishConnectionRun is the reconnect/Stop generation boundary. The second
// send-queue drain is intentionally after the owner join: runPumps drains its
// own carrier output, but a cancellation-aware subordinate may enqueue a final
// response while unwinding. No next connection may observe those old frames.
func (c *Protocol) finishConnectionRun(run *connectionRun, stopSessionTelemetry func(), reason string) {
	if run == nil {
		return
	}
	owners := run.closeAdmission()
	c.closeElasticRequestLinks(reason)
	if stopSessionTelemetry != nil {
		stopSessionTelemetry()
	}
	_ = waitForConnectionRunOwners(context.Background(), owners)
	c.drainSendQueue()
}

// startConnectionRunTask preserves the historical asynchronous behaviour for
// package helpers invoked directly by tests or process-scoped callers, while
// production Session Link contexts always take the exact run admission path.
func (c *Protocol) startConnectionRunTask(ctx context.Context, task func(context.Context)) bool {
	if task == nil {
		return false
	}
	if run := connectionRunFromContext(ctx); run != nil {
		return run.start(ctx, task)
	}
	if ctx == nil {
		ctx = context.Background()
	}
	go task(ctx)
	return true
}