// Package streamregistry_test exercises the public stream registry
// API: typed cancel causes, first-wins Release/Cancel arbitration,
// idempotent CancelAll, and bounded shutdown semantics. The registry
// itself owns NO body-interrupt path — gateway-side read-deadline
// interrupt is verified separately in the api package tests.
package streamregistry_test

import (
	"context"
	"errors"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/tidbcloud/AgentTeam/desktop/internal/streamregistry"
)

const (
	// testEntryCancelTimeout bounds any wait for a typed cancel to
	// propagate. Every test path must observe its terminal state
	// within this window; a hang indicates Cancel failed to fire
	// (or that Release raced ahead and lost when it should not).
	testEntryCancelTimeout = 2 * time.Second
	// testCancelAllBound is the upper bound for the synchronous
	// portion of CancelAll. If CancelAll itself takes longer, it is
	// leaking into per-entry cleanup and the bounded-shutdown
	// property is broken.
	testCancelAllBound = 500 * time.Millisecond
	// testParentDeadlineBudget is the deadline we install on a
	// parent context to verify the deadline bridge. Generous enough
	// to be stable under -race, tight enough that the test does not
	// sleep noticeably.
	testParentDeadlineBudget = 15 * time.Millisecond
)

// waitCtx returns a context bounded by testEntryCancelTimeout.
func waitCtx(t *testing.T) context.Context {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), testEntryCancelTimeout)
	t.Cleanup(cancel)
	return ctx
}

// requireEntryCanceled waits for e.Context().Done() to close and
// asserts the cause satisfies want. want==nil asserts the context
// stayed live (no cause); callers use it to prove normal completion
// did NOT introduce a cancel.
func requireEntryCanceled(t *testing.T, e *streamregistry.Entry, want error) {
	t.Helper()
	select {
	case <-e.Context().Done():
	case <-time.After(testEntryCancelTimeout):
		t.Fatalf("entry %p: Done() never closed within %v", e, testEntryCancelTimeout)
	}
	if want == nil {
		if cause := context.Cause(e.Context()); cause != nil {
			t.Fatalf("entry %p: cause = %v, want nil", e, cause)
		}
		return
	}
	cause := context.Cause(e.Context())
	if !errors.Is(cause, want) {
		t.Fatalf("entry %p: cause = %v, want %v", e, cause, want)
	}
}

// requireEntryStillLive asserts the entry context has not been
// canceled. Callers use it to prove a no-op path (Release-wins,
// sibling session not targeted, etc.) actually left the entry alone.
func requireEntryStillLive(t *testing.T, e *streamregistry.Entry) {
	t.Helper()
	select {
	case <-e.Context().Done():
		t.Fatalf("entry %p: Done() closed unexpectedly; cause=%v", e, context.Cause(e.Context()))
	default:
	}
	if cause := context.Cause(e.Context()); cause != nil {
		t.Fatalf("entry %p: cause = %v, want nil", e, cause)
	}
}

// TestRegistryRegisterAdmitsAndCounts verifies basic register/release
// accounting. After Register the active count is 1; after Release it
// returns to 0; Release is idempotent (a second call does not
// underflow the counter).
func TestRegistryRegisterAdmitsAndCounts(t *testing.T) {
	r := streamregistry.New()
	if got := r.ActiveGatewayOps(); got != 0 {
		t.Fatalf("initial count = %d, want 0", got)
	}
	e, err := r.Register(context.Background(), "s1", 1, "srv-a")
	if err != nil {
		t.Fatalf("Register: %v", err)
	}
	if got := r.ActiveGatewayOps(); got != 1 {
		t.Fatalf("after register = %d, want 1", got)
	}
	e.Release()
	e.Release()
	if got := r.ActiveGatewayOps(); got != 0 {
		t.Fatalf("after release = %d, want 0", got)
	}
}

// TestRegistryRegisterConcurrentSameSessionAndGeneration verifies
// that concurrent registrations with the same (sessionID, generation,
// serverID) get distinct OpTokens and thus distinct Entries. This is
// the invariant that lets a session multiplex multiple concurrent
// gateway ops without one entry overwriting another.
func TestRegistryRegisterConcurrentSameSessionAndGeneration(t *testing.T) {
	r := streamregistry.New()
	const n = 32
	keys := make([]streamregistry.Key, n)
	errs := make([]error, n)
	var wg sync.WaitGroup
	wg.Add(n)
	for i := 0; i < n; i++ {
		i := i
		go func() {
			defer wg.Done()
			e, err := r.Register(context.Background(), "s1", 7, "srv-a")
			errs[i] = err
			if err != nil {
				return
			}
			keys[i] = e.Key()
			defer e.Release()
		}()
	}
	wg.Wait()
	seen := make(map[uint64]bool, n)
	for i := 0; i < n; i++ {
		if errs[i] != nil {
			t.Fatalf("Register %d: %v", i, errs[i])
		}
		k := keys[i]
		if k.SessionID != "s1" || k.Generation != 7 || k.ServerID != "srv-a" {
			t.Fatalf("entry %d unexpected key: %+v", i, k)
		}
		if k.OpToken == 0 || seen[k.OpToken] {
			t.Fatalf("entry %d duplicate/zero OpToken: %d", i, k.OpToken)
		}
		seen[k.OpToken] = true
	}
}

// TestRegistryRegisterFailsFastAfterCancelAll verifies that once
// CancelAll has flipped the registry into draining state, Register
// returns ErrRegistryClosed and does NOT insert a new entry. This
// closes the shutdown race where a late handler could otherwise
// re-arm after CancelAll's snapshot pass.
func TestRegistryRegisterFailsFastAfterCancelAll(t *testing.T) {
	r := streamregistry.New()
	wait := r.CancelAll()
	if err := wait(waitCtx(t)); err != nil {
		t.Fatalf("wait: %v", err)
	}
	if _, err := r.Register(context.Background(), "s1", 1, "srv-a"); !errors.Is(err, streamregistry.ErrRegistryClosed) {
		t.Fatalf("err = %v, want ErrRegistryClosed", err)
	}
	if got := r.ActiveGatewayOps(); got != 0 {
		t.Fatalf("count = %d, want 0 (no entry inserted)", got)
	}
}

// TestRegistryNormalReleaseHasNoCancelCause verifies that normal
// completion via Release does NOT introduce a typed cause. A handler
// that returns EOF/upstream-error must not surface as context.Canceled
// to upstream observers (logs, metrics, response classifier).
func TestRegistryNormalReleaseHasNoCancelCause(t *testing.T) {
	r := streamregistry.New()
	e, _ := r.Register(context.Background(), "s1", 1, "srv-a")
	e.Release()
	requireEntryStillLive(t, e)
}

// TestRegistryCancelAllFiresTypedDaemonShutdown verifies that every
// active entry is canceled with the daemon-shutdown typed cause and
// that the typed CancelError is recoverable via errors.As against
// the bare Reason as well.
func TestRegistryCancelAllFiresTypedDaemonShutdown(t *testing.T) {
	r := streamregistry.New()
	e1, _ := r.Register(context.Background(), "s1", 1, "srv-a")
	e2, _ := r.Register(context.Background(), "s2", 2, "srv-b")
	e3, _ := r.Register(context.Background(), "s1", 2, "srv-a")

	r.CancelAll()(waitCtx(t))

	for i, e := range []*streamregistry.Entry{e1, e2, e3} {
		requireEntryCanceled(t, e, streamregistry.ReasonDaemonShutdown)
		var ce streamregistry.CancelError
		if !errors.As(context.Cause(e.Context()), &ce) {
			t.Fatalf("entry %d: cause is not CancelError: %v", i, context.Cause(e.Context()))
		}
		if ce.Reason != streamregistry.ReasonDaemonShutdown {
			t.Fatalf("entry %d: reason = %q, want %q", i, ce.Reason, streamregistry.ReasonDaemonShutdown)
		}
	}
}

// TestRegistryCancelSessionGenerationTargetsSessionAndGeneration
// verifies CancelSessionGeneration only touches entries matching
// both sessionID and generation. Entries in other sessions, the
// same session at a different generation, and other sessions at
// that generation all remain running.
func TestRegistryCancelSessionGenerationTargetsSessionAndGeneration(t *testing.T) {
	r := streamregistry.New()
	target, _ := r.Register(context.Background(), "s1", 7, "srv-a")
	otherSession, _ := r.Register(context.Background(), "s2", 7, "srv-b")
	otherGen, _ := r.Register(context.Background(), "s1", 8, "srv-a")

	r.CancelSessionGeneration("s1", 7, streamregistry.ReasonGenerationChange)

	requireEntryCanceled(t, target, streamregistry.ReasonGenerationChange)
	requireEntryStillLive(t, otherSession)
	requireEntryStillLive(t, otherGen)
	otherSession.Release()
	otherGen.Release()
}

// TestRegistryParentCancelBecomesClientDisconnect verifies that when
// the parent context (typically the inbound *http.Request context)
// is canceled, the entry is canceled with ReasonClientDisconnect.
// This is the bridge that lets the gateway tear down reads when the
// client goes away without the handler polling.
func TestRegistryParentCancelBecomesClientDisconnect(t *testing.T) {
	r := streamregistry.New()
	parent, cancelParent := context.WithCancel(context.Background())
	e, _ := r.Register(parent, "s1", 1, "srv-a")

	cancelParent()
	requireEntryCanceled(t, e, streamregistry.ReasonClientDisconnect)
	e.Release()
}

// TestRegistryParentDeadlineBecomesDeadlineExceeded verifies parent
// deadline expiry is bridged as context.DeadlineExceeded (NOT a
// typed CancelError). Keeping the cause bare lets observers tell
// "client disconnect" (typed CancelError) apart from "parent timeout"
// (bare DeadlineExceeded) via errors.As/Is without ambiguity. The
// gateway does NOT synthesize a 504 from this cause; the distinction
// exists for log/metric clarity only.
func TestRegistryParentDeadlineBecomesDeadlineExceeded(t *testing.T) {
	r := streamregistry.New()
	parent, cancelParent := context.WithTimeout(context.Background(), testParentDeadlineBudget)
	defer cancelParent()
	e, _ := r.Register(parent, "s1", 1, "srv-a")

	requireEntryCanceled(t, e, context.DeadlineExceeded)
	var ce streamregistry.CancelError
	if errors.As(context.Cause(e.Context()), &ce) {
		t.Fatalf("cause %v is a CancelError; want bare DeadlineExceeded so observers can distinguish parent timeout from typed stream cancel", ce)
	}
	e.Release()
}

// TestRegistryEntryReleaseFirstMakesCancelNoOp verifies the first-
// wins arbitration: if Release wins, the operation context stays
// live (no cause), and a subsequent Cancel call is a no-op (does
// NOT close Done() retroactively). This is the property that lets
// the gateway call Cancel defensively after a successful ResponseWriter
// write without flipping a completed op into a "canceled" state.
func TestRegistryEntryReleaseFirstMakesCancelNoOp(t *testing.T) {
	r := streamregistry.New()
	e, _ := r.Register(context.Background(), "s1", 1, "srv-a")

	e.Release()
	e.Cancel(streamregistry.CancelError{Reason: streamregistry.ReasonDaemonShutdown})
	requireEntryStillLive(t, e)
}

// TestRegistryEntryCancelFirstStableCauseReleaseWaits verifies the
// other half of first-wins: if Cancel wins, the cause is stable and
// a subsequent Release does NOT overwrite it. Release must also wait
// for the cancel side-effect (cancelApplied) so observers can read
// cause immediately after Release returns.
func TestRegistryEntryCancelFirstStableCauseReleaseWaits(t *testing.T) {
	r := streamregistry.New()
	e, _ := r.Register(context.Background(), "s1", 1, "srv-a")
	want := streamregistry.CancelError{Reason: streamregistry.ReasonGenerationChange}

	e.Cancel(want)
	requireEntryCanceled(t, e, streamregistry.ReasonGenerationChange)

	e.Release()
	if cause := context.Cause(e.Context()); !errors.Is(cause, want) {
		t.Fatalf("cause after Release = %v, want %v", cause, want)
	}
}

// TestRegistryCancelIdempotentStableFirstCause verifies that calling
// Cancel twice with different causes does not flip the recorded
// cause. The first cause is sticky.
func TestRegistryCancelIdempotentStableFirstCause(t *testing.T) {
	r := streamregistry.New()
	e, _ := r.Register(context.Background(), "s1", 1, "srv-a")
	first := streamregistry.CancelError{Reason: streamregistry.ReasonDaemonShutdown}
	second := streamregistry.CancelError{Reason: streamregistry.ReasonGenerationChange}

	e.Cancel(first)
	e.Cancel(second)

	requireEntryCanceled(t, e, streamregistry.ReasonDaemonShutdown)
}

// TestRegistryCancelAllIdempotentReturnsSameDrained verifies that
// two CancelAll calls return wait functions that observe the same
// drained channel; the shutdown cause is sticky across the second
// call (the first caller's intent is preserved).
func TestRegistryCancelAllIdempotentReturnsSameDrained(t *testing.T) {
	r := streamregistry.New()
	e, _ := r.Register(context.Background(), "s1", 1, "srv-a")

	wait1 := r.CancelAll()
	wait2 := r.CancelAll()

	requireEntryCanceled(t, e, streamregistry.ReasonDaemonShutdown)

	e.Release()
	if err := wait1(waitCtx(t)); err != nil {
		t.Fatalf("wait1: %v", err)
	}
	if err := wait2(waitCtx(t)); err != nil {
		t.Fatalf("wait2: %v", err)
	}
}

// TestRegistryCancelAllWaitRespectsContextCancel verifies the wait
// function returns ctx.Err() when ctx expires BEFORE drained closes.
// This is what lets daemon shutdown fall back to "still Preparing"
// if the gateway cannot prove every read exited.
func TestRegistryCancelAllWaitRespectsContextCancel(t *testing.T) {
	r := streamregistry.New()
	// Entry never released — drained cannot close.
	e, _ := r.Register(context.Background(), "s1", 1, "srv-a")

	wait := r.CancelAll()

	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Millisecond)
	defer cancel()
	if err := wait(ctx); !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("wait = %v, want %v", err, context.DeadlineExceeded)
	}
	// Typed cancel fired even though wait timed out.
	requireEntryCanceled(t, e, streamregistry.ReasonDaemonShutdown)
}

// TestRegistryCancelAllSynchronousPortionIsBounded is the key
// anti-blocking test: even when no entry is ever Released, the
// synchronous portion of CancelAll (snapshot + Cancel fanout) MUST
// return within a tight bound. The registry MUST NOT accept any
// external Closer/interrupt that could make an entry's Cancel block
// on a wedged upstream; the wait function is the only place a caller
// may block, and it is bounded by the caller's ctx.
func TestRegistryCancelAllSynchronousPortionIsBounded(t *testing.T) {
	r := streamregistry.New()
	for i := 0; i < 16; i++ {
		_, _ = r.Register(context.Background(), "s1", uint64(i+1), "srv-a")
	}
	start := time.Now()
	_ = r.CancelAll()
	if elapsed := time.Since(start); elapsed > testCancelAllBound {
		t.Fatalf("CancelAll synchronous portion took %v, want ≤ %v", elapsed, testCancelAllBound)
	}
}

// TestRegistryConcurrentReleaseVsCancelNoPanic verifies the
// first-wins arbitration is race-free under -race. Multiple goroutines
// race Release and Cancel against distinct entries; the test asserts
// no panic, no double-close, every Released channel closes exactly
// once, and count returns to zero.
func TestRegistryConcurrentReleaseVsCancelNoPanic(t *testing.T) {
	r := streamregistry.New()
	const n = 200
	var wg sync.WaitGroup
	wg.Add(n)
	var observed atomic.Int64
	for i := 0; i < n; i++ {
		i := i
		go func() {
			defer wg.Done()
			e, err := r.Register(context.Background(), "s1", uint64(i+1), "srv-a")
			if err != nil {
				t.Errorf("Register %d: %v", i, err)
				return
			}
			var inner sync.WaitGroup
			inner.Add(2)
			go func() {
				defer inner.Done()
				e.Cancel(streamregistry.CancelError{Reason: streamregistry.ReasonDaemonShutdown})
			}()
			go func() {
				defer inner.Done()
				e.Release()
			}()
			inner.Wait()
			<-e.Released()
			observed.Add(1)
		}()
	}
	wg.Wait()
	if got := observed.Load(); got != n {
		t.Fatalf("observed %d completed Released channels, want %d", got, n)
	}
	if got := r.ActiveGatewayOps(); got != 0 {
		t.Fatalf("ActiveGatewayOps = %d, want 0 after all races settle", got)
	}
}

// TestRegistryReleasedStopsParentBridge verifies that after Release,
// the AfterFunc bridge installed on the parent context is STOPPED — a
// subsequent parent cancellation must NOT invoke the bridge callback
// and must NOT flip the entry's cause. Without this guarantee, a
// normal-completed operation could be retroactively marked as a
// client disconnect when the underlying request context fires its
// late cancellation chain (e.g. deferred ctx cancel after handler
// return).
type recordingAfterFuncContext struct {
	context.Context
	mu        sync.Mutex
	callback  func()
	installed chan struct{}
	stopped   chan struct{}
	stopOnce  sync.Once
}

func newRecordingAfterFuncContext() *recordingAfterFuncContext {
	return &recordingAfterFuncContext{
		Context:   context.Background(),
		installed: make(chan struct{}),
		stopped:   make(chan struct{}),
	}
}

func (c *recordingAfterFuncContext) AfterFunc(callback func()) func() bool {
	c.mu.Lock()
	c.callback = callback
	c.mu.Unlock()
	close(c.installed)
	return func() bool {
		c.stopOnce.Do(func() { close(c.stopped) })
		return true
	}
}

func (c *recordingAfterFuncContext) Done() <-chan struct{} {
	// Non-nil forces context.AfterFunc to register through this context's
	// AfterFunc method rather than treating it as a never-canceled parent.
	return make(chan struct{})
}

// Evidence: the custom parent records the real context.AfterFunc registration
// and the exact stop function invocation. This distinguishes "Release-first
// arbitration ignored a late callback" from "Release actually unregistered
// the bridge".
func TestRegistryReleasedStopsParentBridge(t *testing.T) {
	r := streamregistry.New()
	parent := newRecordingAfterFuncContext()
	e, err := r.Register(parent, "s1", 1, "srv-a")
	if err != nil {
		t.Fatalf("Register: %v", err)
	}
	select {
	case <-parent.installed:
	case <-time.After(testEntryCancelTimeout):
		t.Fatal("parent AfterFunc bridge was not installed")
	}

	e.Release()
	requireEntryStillLive(t, e)
	select {
	case <-parent.stopped:
	case <-time.After(testEntryCancelTimeout):
		t.Fatal("Release did not call the real parent AfterFunc stop function")
	}
	if cause := context.Cause(e.Context()); cause != nil {
		t.Fatalf("cause after Release = %v, want nil", cause)
	}
}

// TestRegistryReleasedStopsDeadlineTimer verifies the deadline bridge
// (time.AfterFunc variant) is also stopped on Release. Otherwise an
// already-completed operation whose parent had a near deadline could
// retroactively flip to context.DeadlineExceeded, masking normal
// completion in observability.
func TestRegistryReleasedStopsDeadlineTimer(t *testing.T) {
	r := streamregistry.New()
	parent, cancelParent := context.WithTimeout(context.Background(), testParentDeadlineBudget)
	defer cancelParent()

	e, err := r.Register(parent, "s1", 1, "srv-a")
	if err != nil {
		t.Fatalf("Register: %v", err)
	}

	e.Release()
	requireEntryStillLive(t, e)

	// Wait beyond the original deadline; without a timer Stop, the
	// AfterFunc would fire and flip cause to DeadlineExceeded.
	time.Sleep(testParentDeadlineBudget * 3)

	if cause := context.Cause(e.Context()); cause != nil {
		t.Fatalf("cause flipped after Release + deadline expiry: %v (want nil — Release must stop the deadline timer)", cause)
	}
	select {
	case <-e.Context().Done():
		t.Fatalf("entry Context().Done closed after Release + deadline expiry")
	default:
	}
}

// TestRegistryCancelAllVsReleaseSameEntryFirstCauseStable verifies the
// FIRST-CAUSE stability contract on a SINGLE entry under race. Two
// goroutines fire CancelAll (which would cancel with ReasonDaemonShutdown)
// and entry.Release simultaneously; whichever wins, the post-state MUST
// satisfy:
//
//   - If Release won: entry.Context() stays live (cause==nil), entry is
//     released exactly once (Released closed once), activeOps returns
//     to zero, drained closes promptly. CancelAll's wait function
//     returns nil because count==0.
//   - If CancelAll won: entry.Context() is canceled with
//     ReasonDaemonShutdown (the daemon-shutdown typed cause), the cause
//     is sticky across the subsequent Release (Release must NOT
//     overwrite), and activeOps returns to zero.
//
// Either outcome is acceptable; the test asserts the OR of both
// post-states. The contract being tested: the registry's first-wins
// arbitration must produce a STABLE post-state — never "released but
// cause flipped", never "canceled twice with different causes", never
// "activeOps leaks past Release", never "drained never closes".
func TestRegistryCancelAllVsReleaseSameEntryFirstCauseStable(t *testing.T) {
	const iterations = 200
	for iter := 0; iter < iterations; iter++ {
		r := streamregistry.New()
		e, err := r.Register(context.Background(), "s1", 1, "srv-a")
		if err != nil {
			t.Fatalf("iter %d: Register: %v", iter, err)
		}

		var wg sync.WaitGroup
		wg.Add(2)
		var waitFn func(context.Context) error
		go func() {
			defer wg.Done()
			waitFn = r.CancelAll()
		}()
		go func() {
			defer wg.Done()
			e.Release()
		}()
		wg.Wait()

		// Released MUST close exactly once (releaseOnce guard).
		select {
		case <-e.Released():
		case <-time.After(testEntryCancelTimeout):
			t.Fatalf("iter %d: Released never closed", iter)
		}

		// Cause is binary: nil (Release won) or ReasonDaemonShutdown
		// (CancelAll won). Any other cause is a first-cause violation.
		cause := context.Cause(e.Context())
		switch {
		case cause == nil:
			// Release won: entry context must stay live.
			select {
			case <-e.Context().Done():
				t.Fatalf("iter %d: Release won but Context().Done closed", iter)
			default:
			}
		case errors.Is(cause, streamregistry.ReasonDaemonShutdown):
			// CancelAll won: entry context must be canceled.
			select {
			case <-e.Context().Done():
			default:
				t.Fatalf("iter %d: CancelAll won but Context().Done still open", iter)
			}
		default:
			t.Fatalf("iter %d: cause = %v, want nil or ReasonDaemonShutdown", iter, cause)
		}

		// activeOps MUST return to zero regardless of winner.
		if got := r.ActiveGatewayOps(); got != 0 {
			t.Fatalf("iter %d: ActiveGatewayOps = %d, want 0", iter, got)
		}

		// CancelAll's wait function MUST return nil because count==0
		// and closed==true → drained closes immediately. If waitFn
		// times out, drained is not wired up correctly.
		if waitFn == nil {
			t.Fatalf("iter %d: CancelAll did not return a wait function", iter)
		}
		ctx, cancel := context.WithTimeout(context.Background(), testEntryCancelTimeout)
		defer cancel()
		if err := waitFn(ctx); err != nil {
			t.Fatalf("iter %d: CancelAll wait: %v (drained must close when count==0)", iter, err)
		}

		// Sticky-cause re-check: a second Cancel (direct, with a
		// different reason) MUST NOT flip the cause. This proves the
		// first-cause is durable even after the registry-level race.
		e.Cancel(streamregistry.CancelError{Reason: streamregistry.ReasonGenerationChange})
		if afterCause := context.Cause(e.Context()); afterCause != cause {
			t.Fatalf("iter %d: cause flipped after second Cancel: got %v, want %v (first-cause must be sticky)",
				iter, afterCause, cause)
		}
	}
}