// Package streamregistry tracks in-flight remote gateway operations so
// daemon-wide lifecycle events (server switch, shutdown, client
// disconnect) can cancel them with a typed, observable cause.
package streamregistry

import (
	"context"
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

// Reason is a typed stream-cancellation cause.
type Reason string

const (
	ReasonGenerationChange Reason = "generation-change"
	ReasonDaemonShutdown   Reason = "daemon-shutdown"
	ReasonClientDisconnect Reason = "client-disconnect"
)

func (r Reason) Error() string { return "stream canceled: " + string(r) }

// CancelError wraps a Reason so it can be passed to
// context.CancelCauseFunc and recovered via errors.Is against either
// the Reason or the wrapper.
type CancelError struct{ Reason Reason }

func (e CancelError) Error() string { return "stream canceled: " + string(e.Reason) }
func (e CancelError) Unwrap() error { return e.Reason }

// ErrRegistryClosed is returned by Register after CancelAll has flipped
// the registry into the draining state. The cause is fixed to
// daemon-shutdown; the first CancelAll caller's intent is preserved
// across repeat calls.
var ErrRegistryClosed = CancelError{Reason: ReasonDaemonShutdown}

// Key uniquely identifies an in-flight remote gateway operation. OpToken
// disambiguates concurrent ops within the same (session, generation).
type Key struct {
	SessionID  string
	Generation uint64
	ServerID   string
	OpToken    uint64
}

type entryState int

const (
	entryIdle entryState = iota
	entryCommitted
	entryReleased
	entryCanceled
)

// deadlineValueCtx preserves parent values and Deadline metadata without
// inheriting parent Done/Err. Used as the base for Entry.Context() so
// parent cancellation does NOT propagate — only an explicit Entry.Cancel
// call cancels the operation context, with a typed cause.
type deadlineValueCtx struct {
	context.Context
	deadline time.Time
	hasDL    bool
}

func (c *deadlineValueCtx) Deadline() (time.Time, bool) { return c.deadline, c.hasDL }
func (c *deadlineValueCtx) Done() <-chan struct{}       { return nil }
func (c *deadlineValueCtx) Err() error                  { return nil }

// Entry tracks one in-flight op. Holders call Release on normal
// completion (EOF, upstream error). Holders call Cancel to terminate
// with a typed cause; the cause is observable via context.Cause on
// Context().
//
// Release and Cancel race are arbitrated first-wins:
//   - Release first → state goes idle→released; cause stays nil;
//     subsequent Cancel is a no-op.
//   - Cancel first → state goes idle→canceled; first cause is stable;
//     subsequent Cancel is a no-op; Release runs cleanup but does NOT
//     overwrite the cause.
//
// When Release observes state==canceled, it waits on cancelApplied so
// the cancel side-effect (cancelCause) has completed before resources
// are released. This lets CancelAll's bounded wait converge: by the
// time activeOps reaches zero, every entry's typed cancel has actually
// fired. The registry deliberately does NOT accept an external body
// interrupt callback: a server-side *http.body.Close blocks on the
// read mutex (and may drain up to 256KiB before returning), so wiring
// it into Entry.Cancel would let a slow upstream pin CancelAll. The
// gateway instead interrupts reads via ResponseController.SetReadDeadline
// on entry.Context().Done(), keeping the registry wait bounded.
type Entry struct {
	registry *Registry
	key      Key

	opCtx  context.Context
	cancel context.CancelCauseFunc

	stateMu        sync.Mutex
	state          entryState
	cause          error
	cancelApplied  chan struct{}
	released       chan struct{}
	commitObserved chan struct{}

	deadlineTimer *time.Timer
	stopBridge    func() bool
	releaseOnce   sync.Once
}

// Key returns the registry-assigned identifier.
func (e *Entry) Key() Key { return e.key }

// Context returns the operation context. It is canceled (with cause)
// on Entry.Cancel; it is NOT canceled on parent cancel unless that
// cancel is bridged via Entry.Cancel internally. Deadline metadata
// from the parent is preserved.
func (e *Entry) Context() context.Context { return e.opCtx }

// Released returns a channel closed when Release completes (normal or
// after a cancel). Watchers and bridges select on this to exit on
// normal completion without polling Context().Done().
func (e *Entry) Released() <-chan struct{} { return e.released }

// BindInterrupt is intentionally NOT provided. The registry refuses
// to accept any external body-interrupt callback: a server-side
// *http.body.Close blocks on the read mutex and may drain up to 256KiB
// before returning, so wiring it through Entry.Cancel would let a slow
// upstream pin CancelAll indefinitely. The gateway must instead
// interrupt in-flight reads on entry.Context().Done() via
// http.NewResponseController(w).SetReadDeadline(time.Now()), which
// acts on the underlying connection and is non-blocking; the handler
// then closes the body and calls Release from the same goroutine that
// owns the read.

// TryCommit is the single owner-transfer point from the pre-commit
// registry-owned phase to the post-commit pump-owned phase. Production
// callers (the WS gateway) invoke it after the upstream 101 + subprotocol
// verification but BEFORE the downstream Hijack/Upgrade commit. It is the
// handoff that decides whether the gateway may commit the downstream
// connection.
//
// TryCommit is arbitrated against Cancel first-wins under stateMu. The
// returned (owned, firstCause) tuple encodes the arbitration result:
//   - (true, nil): state transitioned idle→committed. The caller owns the
//     commit step and may proceed to Hijack/Upgrade. A subsequent Cancel
//     still fires (committed→canceled) so post-commit pumps observe
//     entry.Context().Done() and close with typed close codes.
//   - (false, cause): Cancel won first (state is canceled). cause is the
//     stored Entry cause — the gateway surfaces the typed HTTP contract
//     (409 for generation change, 503 for daemon shutdown, no-write for
//     client disconnect / deadline) and MUST NOT proceed to the commit.
//     cause is read atomically from Entry.cause so TryCommit cannot
//     observe (false, nil) when Entry.Context() is already canceled.
//   - (false, nil): state is released OR already committed. The entry is
//     either gone or has already transferred commit ownership to a
//     previous caller; no (re-)commit is possible.
//
// Repeated TryCommit never re-grants commit ownership: a second call on an
// already-committed entry returns (false, nil) without re-closing the
// CommitObserved signal. A second call on a since-canceled entry returns
// (false, original-cause) via the canceled branch.
func (e *Entry) TryCommit() (owned bool, firstCause error) {
	e.stateMu.Lock()
	switch e.state {
	case entryCommitted:
		e.stateMu.Unlock()
		return false, nil
	case entryReleased:
		e.stateMu.Unlock()
		return false, nil
	case entryCanceled:
		cause := e.cause
		e.stateMu.Unlock()
		return false, cause
	case entryIdle:
		e.state = entryCommitted
		close(e.commitObserved)
		e.stateMu.Unlock()
		return true, nil
	}
	e.stateMu.Unlock()
	return false, nil
}

// CommitObserved returns a channel that is closed exactly once when
// TryCommit successfully acquires commit ownership (idle→committed).
// Observers select on this to synchronize with the commit-owned boundary
// without influencing TryCommit's arbitration semantics. The channel is
// never closed if TryCommit loses to Cancel/Release. It is a read-only
// observer seam: external code cannot inject behavior through it.
func (e *Entry) CommitObserved() <-chan struct{} { return e.commitObserved }

// Cancel terminates the operation with cause. First-wins; subsequent
// Cancel calls are no-ops. cancelCause runs synchronously;
// cancelApplied is closed once it finishes so any Release observing
// canceled can proceed without racing the side-effect.
//
// A nil cause is normalised to context.Canceled so TryCommit's
// firstCause and context.Cause(entry.Context()) report a stable,
// non-nil value — callers never need to handle (false, nil) on a
// canceled entry.
//
// Cancel from the committed state still fires: post-commit pumps rely
// on entry.Context().Done() to close with typed close codes (1012 for
// generation change, 1001 for daemon shutdown, etc.).
func (e *Entry) Cancel(cause error) {
	e.stateMu.Lock()
	switch e.state {
	case entryReleased, entryCanceled:
		e.stateMu.Unlock()
		return
	}
	if cause == nil {
		cause = context.Canceled
	}
	e.state = entryCanceled
	e.cause = cause
	e.stateMu.Unlock()

	e.cancel(cause)
	close(e.cancelApplied)
}

// Release is idempotent. On normal completion (state idle/committed →
// released) it does NOT introduce a cancel cause. If Cancel won first,
// Release waits on cancelApplied then runs cleanup without overwriting
// cause.
func (e *Entry) Release() {
	e.stateMu.Lock()
	if e.state == entryReleased {
		e.stateMu.Unlock()
		return
	}
	prev := e.state
	if prev == entryIdle || prev == entryCommitted {
		e.state = entryReleased
	}
	e.stateMu.Unlock()

	if prev == entryCanceled {
		<-e.cancelApplied
	}
	e.releaseOnce.Do(func() {
		if e.stopBridge != nil {
			e.stopBridge()
		}
		if e.deadlineTimer != nil {
			e.deadlineTimer.Stop()
		}
		if e.registry != nil {
			e.registry.releaseEntry(e)
		}
		close(e.released)
	})
}

// Registry tracks in-flight gateway ops and arbitrates daemon-wide
// cancellation. A single drained channel is closed once (closed &&
// count==0); CancelAll is idempotent and returns the same wait
// function with a stable first cause.
type Registry struct {
	mu               sync.Mutex
	closed           bool
	shutdownCause    error
	count            int
	drained          chan struct{}
	drainedOnce      sync.Once
	entries          map[Key]*Entry
	opTokens         map[string]uint64
	activeGatewayOps atomic.Int64
}

// New constructs an empty Registry.
func New() *Registry {
	return &Registry{
		drained:  make(chan struct{}),
		entries:  make(map[Key]*Entry),
		opTokens: make(map[string]uint64),
	}
}

// Register attempts to admit a new op. After CancelAll, Register
// fails fast with ErrRegistryClosed and never inserts an entry.
//
// The returned Entry has a context derived from a deadline-preserving
// base (parent values + Deadline metadata, no Done/Err inheritance).
// Parent cancellation is bridged via context.AfterFunc into
// Entry.Cancel with ReasonClientDisconnect; parent deadline expiry is
// bridged via a stoppable timer into Entry.Cancel(context.DeadlineExceeded).
func (r *Registry) Register(parent context.Context, sessionID string, generation uint64, serverID string) (*Entry, error) {
	if parent == nil {
		parent = context.Background()
	}
	r.mu.Lock()
	if r.closed {
		r.mu.Unlock()
		return nil, ErrRegistryClosed
	}
	tokenKey := sessionID + "\x00" + strconv.FormatUint(generation, 10)
	r.opTokens[tokenKey]++
	token := r.opTokens[tokenKey]
	key := Key{SessionID: sessionID, Generation: generation, ServerID: serverID, OpToken: token}

	base := context.WithoutCancel(parent)
	var deadlineTimer *time.Timer
	if dl, ok := parent.Deadline(); ok {
		base = &deadlineValueCtx{Context: base, deadline: dl, hasDL: true}
		delay := time.Until(dl)
		// placeholder; real timer assigned after entry exists
		_ = delay
	}
	opCtx, cancel := context.WithCancelCause(base)

	e := &Entry{
		registry:       r,
		key:            key,
		opCtx:          opCtx,
		cancel:         cancel,
		cancelApplied:  make(chan struct{}),
		released:       make(chan struct{}),
		commitObserved: make(chan struct{}),
	}
	e.stopBridge = context.AfterFunc(parent, func() {
		if parent.Err() == context.DeadlineExceeded {
			return
		}
		e.Cancel(CancelError{Reason: ReasonClientDisconnect})
	})
	if dl, ok := parent.Deadline(); ok {
		delay := time.Until(dl)
		deadlineTimer = time.AfterFunc(delay, func() {
			e.Cancel(context.DeadlineExceeded)
		})
		e.deadlineTimer = deadlineTimer
	}

	r.entries[key] = e
	r.count++
	r.activeGatewayOps.Add(1)
	r.mu.Unlock()
	return e, nil
}

// CancelSessionGeneration cancels every entry matching sessionID+
// generation with the given reason. Snapshot is taken under the lock;
// actual cancel runs outside to avoid callback reentry.
func (r *Registry) CancelSessionGeneration(sessionID string, generation uint64, reason Reason) {
	r.mu.Lock()
	snapshot := make([]*Entry, 0)
	for k, e := range r.entries {
		if k.SessionID == sessionID && k.Generation == generation {
			snapshot = append(snapshot, e)
		}
	}
	r.mu.Unlock()
	cause := CancelError{Reason: reason}
	for _, e := range snapshot {
		e.Cancel(cause)
	}
}

// CancelAll flips the registry into draining state, cancels every
// active entry with ReasonDaemonShutdown, and returns a wait function
// that blocks until count reaches zero or ctx expires. Repeat calls
// return an equivalent wait function and do NOT overwrite the first
// cause.
func (r *Registry) CancelAll() func(ctx context.Context) error {
	r.mu.Lock()
	if r.closed {
		ch := r.drained
		r.mu.Unlock()
		return func(ctx context.Context) error {
			select {
			case <-ch:
				return nil
			case <-ctx.Done():
				return ctx.Err()
			}
		}
	}
	r.closed = true
	r.shutdownCause = ErrRegistryClosed
	snapshot := make([]*Entry, 0, len(r.entries))
	for _, e := range r.entries {
		snapshot = append(snapshot, e)
	}
	count := r.count
	r.mu.Unlock()

	cause := r.shutdownCause
	for _, e := range snapshot {
		e.Cancel(cause)
	}
	if count == 0 {
		r.drainedOnce.Do(func() { close(r.drained) })
	}
	return func(ctx context.Context) error {
		select {
		case <-r.drained:
			return nil
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// ActiveGatewayOps returns the current count of in-flight entries.
func (r *Registry) ActiveGatewayOps() int64 {
	return r.activeGatewayOps.Load()
}

func (r *Registry) releaseEntry(e *Entry) {
	r.mu.Lock()
	if _, ok := r.entries[e.key]; ok {
		delete(r.entries, e.key)
		r.count--
		if r.count < 0 {
			r.count = 0
		}
		r.activeGatewayOps.Add(-1)
	}
	closed := r.closed
	count := r.count
	r.mu.Unlock()

	if closed && count == 0 {
		r.drainedOnce.Do(func() { close(r.drained) })
	}
}