package api

import (
	"bufio"
	"context"
	"encoding/base64"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"

	"github.com/tidbcloud/AgentTeam/desktop/internal/api/localsecurity"
	"github.com/tidbcloud/AgentTeam/desktop/internal/streamregistry"
	"github.com/tidbcloud/AgentTeam/desktop/internal/types"
	"github.com/tidbcloud/AgentTeam/desktop/internal/uisession"
)

// hijackCapturer is a thin http.ResponseWriter wrapper that intercepts
// gorilla's internal Hijack() call during Upgrade so the pre-upgrade
// cancel-watcher can close the raw downstream conn while the 101
// response Write is still blocked. Every interface method other than
// Hijack is delegated to the embedded ResponseWriter, so gorilla's
// Upgrader (which asserts io.ReaderFrom, http.Flusher, http.Hijacker,
// etc.) sees the same concrete underlying writer it would without the
// wrapper.
//
// Close discipline: closeRawConn records a "close-at-Hijack" flag if
// Hijack has not yet been called; otherwise it closes the captured raw
// conn directly. Either way, the watcher never races the goroutine that
// is about to call (or has just called) Hijack — the mutex serializes
// the capture and the close request.
type hijackCapturer struct {
	http.ResponseWriter
	rawConnMu sync.Mutex
	rawConn   net.Conn
	closeFlag bool
}

// Hijack delegates to the underlying ResponseWriter, captures the raw
// net.Conn for the pre-upgrade watcher, and closes the conn immediately
// if the watcher already requested a close before Hijack landed.
func (h *hijackCapturer) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	conn, brw, err := h.ResponseWriter.(http.Hijacker).Hijack()
	if err != nil {
		return nil, nil, err
	}
	h.rawConnMu.Lock()
	h.rawConn = conn
	closeNow := h.closeFlag
	h.rawConnMu.Unlock()
	if closeNow {
		_ = conn.Close()
	}
	return conn, brw, nil
}

// closeRawConn closes the captured raw conn if Hijack has already been
// called; otherwise marks the capturer to close the conn at Hijack time.
// This lets the pre-upgrade watcher interrupt a blocked 101 Write
// regardless of whether gorilla has reached Hijack yet.
func (h *hijackCapturer) closeRawConn() {
	h.rawConnMu.Lock()
	h.closeFlag = true
	conn := h.rawConn
	h.rawConnMu.Unlock()
	if conn != nil {
		_ = conn.Close()
	}
}

// websocketPreflightError carries the typed error code that P3 maps to the
// pre-admission HTTP contract. INVALID_WEBSOCKET_HANDSHAKE surfaces a 400
// before session/server/registry admission; WEBSOCKET_ORIGIN_REJECTED
// surfaces a 403 at the same boundary. Both are emitted with zero side
// effects: no token fetch, no Registry.Register, no upstream dial.
type websocketPreflightError struct {
	code   string
	detail string
}

func (e websocketPreflightError) Error() string { return e.detail }

// websocketIntentPath is the legacy path prefix the P3 gateway treats as a
// WebSocket endpoint. The exact path (/socket) and its subtree (/socket/...)
// are both admitted as WS intents; sibling paths that merely share the
// literal prefix (e.g. /socketfoo) remain HTTP.
const websocketIntentPath = "/socket"

// isWebSocketIntentPath reports whether the request path is a known upstream WS
// endpoint. The gateway stays on an explicit route set so malformed WS markers
// on unrelated HTTP routes do not change HTTP semantics.
func isWebSocketIntentPath(path string) bool {
	if path == websocketIntentPath {
		return true
	}
	if strings.HasPrefix(path, websocketIntentPath+"/") {
		return true
	}
	return workspaceEventsStreamWorkspaceID(path) != ""
}

// isWebSocketIntent reports whether the request carries any marker that
// identifies it as a WebSocket handshake attempt under a known WS endpoint. A
// single marker (e.g., a bare Sec-WebSocket-Key with no Upgrade, or a
// half-upgrade) is sufficient to enter preflight; preflight itself then rejects
// the malformed intent with INVALID_WEBSOCKET_HANDSHAKE before any
// session/server/registry/dial side effect.
func isWebSocketIntent(request *http.Request) bool {
	if request == nil || request.URL == nil || !isWebSocketIntentPath(request.URL.Path) {
		return false
	}
	header := request.Header
	if header == nil {
		return false
	}
	if _, ok := header["Upgrade"]; ok {
		return true
	}
	if hasConnectionUpgradeToken(header) {
		return true
	}
	for _, marker := range []string{
		"Sec-Websocket-Key",
		"Sec-Websocket-Version",
		"Sec-Websocket-Protocol",
		"Sec-Websocket-Extensions",
	} {
		if _, ok := header[marker]; ok {
			return true
		}
	}
	return false
}

// validateWebSocketHandshake enforces RFC 6455 §4.1 client-opening-handshake
// constraints at the gateway boundary. Any violation returns a typed
// websocketPreflightError so the caller can emit a 400 with code
// INVALID_WEBSOCKET_HANDSHAKE without falling through to the HTTP path.
func validateWebSocketHandshake(request *http.Request) error {
	if request == nil {
		return websocketPreflightError{code: "INVALID_WEBSOCKET_HANDSHAKE", detail: "nil request"}
	}
	if request.Method != http.MethodGet {
		return websocketPreflightError{
			code:   "INVALID_WEBSOCKET_HANDSHAKE",
			detail: fmt.Sprintf("WebSocket handshake must use GET; got %s", request.Method),
		}
	}
	if !hasConnectionUpgradeToken(request.Header) {
		return websocketPreflightError{
			code:   "INVALID_WEBSOCKET_HANDSHAKE",
			detail: "Connection header must list the 'upgrade' token",
		}
	}
	upgrade := request.Header.Values("Upgrade")
	if len(upgrade) != 1 {
		return websocketPreflightError{
			code:   "INVALID_WEBSOCKET_HANDSHAKE",
			detail: "Upgrade header must appear exactly once",
		}
	}
	if !strings.EqualFold(strings.TrimSpace(upgrade[0]), "websocket") {
		return websocketPreflightError{
			code:   "INVALID_WEBSOCKET_HANDSHAKE",
			detail: "Upgrade header must be exactly 'websocket'",
		}
	}
	version := request.Header.Values("Sec-Websocket-Version")
	if len(version) != 1 {
		return websocketPreflightError{
			code:   "INVALID_WEBSOCKET_HANDSHAKE",
			detail: "Sec-WebSocket-Version must appear exactly once",
		}
	}
	if strings.TrimSpace(version[0]) != "13" {
		return websocketPreflightError{
			code:   "INVALID_WEBSOCKET_HANDSHAKE",
			detail: "Sec-WebSocket-Version must be exactly '13'",
		}
	}
	key := request.Header.Values("Sec-Websocket-Key")
	if len(key) != 1 {
		return websocketPreflightError{
			code:   "INVALID_WEBSOCKET_HANDSHAKE",
			detail: "Sec-WebSocket-Key must appear exactly once",
		}
	}
	if !isValidWebSocketKey(key[0]) {
		return websocketPreflightError{
			code:   "INVALID_WEBSOCKET_HANDSHAKE",
			detail: "Sec-WebSocket-Key must be base64-encoded 16 bytes",
		}
	}
	if err := validateWebSocketProtocolHeader(request.Header.Values("Sec-Websocket-Protocol")); err != nil {
		return err
	}
	if err := validateWebSocketExtensionsHeader(request.Header.Values("Sec-Websocket-Extensions")); err != nil {
		return err
	}
	return nil
}

// hasConnectionUpgradeToken reports whether the Connection header (a
// comma-separated token list per RFC 7230 §6.1) contains an "upgrade" token,
// compared case-insensitively. Whitespace inside a token — e.g., "up grade"
// — fails the check because it is not the literal token "upgrade".
func hasConnectionUpgradeToken(header http.Header) bool {
	for _, value := range header.Values("Connection") {
		for _, token := range strings.Split(value, ",") {
			if strings.EqualFold(strings.TrimSpace(token), "upgrade") {
				return true
			}
		}
	}
	return false
}

// isValidWebSocketKey reports whether raw decodes as standard base64 to
// exactly 16 bytes, matching RFC 6455 §4.1 §4 (16-byte randomness).
func isValidWebSocketKey(raw string) bool {
	decoded, err := base64.StdEncoding.DecodeString(raw)
	if err != nil {
		return false
	}
	return len(decoded) == 16
}

// validateWebSocketProtocolHeader enforces that every listed subprotocol is a
// syntactically valid token (RFC 7230 §3.2.6) and that no token appears more
// than once, compared case-insensitively. The header is optional.
func validateWebSocketProtocolHeader(values []string) error {
	if len(values) == 0 {
		return nil
	}
	seen := make(map[string]struct{})
	for _, value := range values {
		for _, token := range strings.Split(value, ",") {
			token = strings.TrimSpace(token)
			if token == "" {
				return websocketPreflightError{
					code:   "INVALID_WEBSOCKET_HANDSHAKE",
					detail: "Sec-WebSocket-Protocol must not contain empty tokens",
				}
			}
			if !isValidHTTPToken(token) {
				return websocketPreflightError{
					code:   "INVALID_WEBSOCKET_HANDSHAKE",
					detail: fmt.Sprintf("Sec-WebSocket-Protocol token %q is not a valid HTTP token", token),
				}
			}
			key := strings.ToLower(token)
			if _, duplicate := seen[key]; duplicate {
				return websocketPreflightError{
					code:   "INVALID_WEBSOCKET_HANDSHAKE",
					detail: fmt.Sprintf("Sec-WebSocket-Protocol token %q appears more than once", token),
				}
			}
			seen[key] = struct{}{}
		}
	}
	return nil
}

// validateWebSocketExtensionsHeader enforces a minimal uniqueness constraint
// on extension names: no extension name may repeat. Offer parsing follows
// RFC 6455 §9.1 (name followed by optional '; params). This Phase 1
// validator rejects syntactically suspicious values but does not validate
// per-extension parameter grammar — full grammar coverage is not exercised
// by the red suite.
func validateWebSocketExtensionsHeader(values []string) error {
	if len(values) == 0 {
		return nil
	}
	seen := make(map[string]struct{})
	for _, value := range values {
		for _, offer := range strings.Split(value, ",") {
			offer = strings.TrimSpace(offer)
			if offer == "" {
				return websocketPreflightError{
					code:   "INVALID_WEBSOCKET_HANDSHAKE",
					detail: "Sec-WebSocket-Extensions must not contain empty offers",
				}
			}
			name := offer
			if idx := strings.IndexByte(name, ';'); idx >= 0 {
				name = strings.TrimSpace(name[:idx])
			}
			if name == "" || !isValidHTTPToken(name) {
				return websocketPreflightError{
					code:   "INVALID_WEBSOCKET_HANDSHAKE",
					detail: fmt.Sprintf("Sec-WebSocket-Extensions offer %q has an invalid name", offer),
				}
			}
			key := strings.ToLower(name)
			if _, duplicate := seen[key]; duplicate {
				return websocketPreflightError{
					code:   "INVALID_WEBSOCKET_HANDSHAKE",
					detail: fmt.Sprintf("Sec-WebSocket-Extensions offer %q appears more than once", name),
				}
			}
			seen[key] = struct{}{}
		}
	}
	return nil
}

// isValidHTTPToken reports whether s is a valid RFC 7230 §3.2.6 token —
// printable US-ASCII excluding delimiters "(),/:;<=>?@[\]{} and space.
func isValidHTTPToken(s string) bool {
	if s == "" {
		return false
	}
	for i := 0; i < len(s); i++ {
		c := s[i]
		if c <= 0x20 || c >= 0x7f {
			return false
		}
		switch c {
		case '(', ')', ',', '/', ':', ';', '<', '=', '>', '?', '@', '[', '\\', ']', '{', '}':
			return false
		}
	}
	return true
}

// validateWebSocketOrigin enforces a strict same-origin policy. Missing
// Origin is allowed (non-browser clients). Otherwise the Origin must:
//   - be a single absolute http(s) URL with no userinfo, opaque, path,
//     query, or fragment;
//   - match the request's scheme (TLS => https, otherwise http) and the
//     request's Host header (case-insensitive).
//
// Any deviation returns a typed websocketPreflightError so the caller can
// emit a 403 with code WEBSOCKET_ORIGIN_REJECTED before admission.
func validateWebSocketOrigin(request *http.Request) error {
	if request == nil {
		return websocketPreflightError{code: "WEBSOCKET_ORIGIN_REJECTED", detail: "nil request"}
	}
	if err := localsecurity.ValidateWebSocketOrigin(request); err != nil {
		return websocketPreflightError{
			code:   "WEBSOCKET_ORIGIN_REJECTED",
			detail: "Origin is not allowed by the local daemon policy",
		}
	}
	return nil
}

// websocketPreflight runs the P3 preflight validators in order: handshake
// constraints first (typed 400), then strict same-origin policy (typed 403).
// The returned error is always a websocketPreflightError so the caller can
// map code/detail to writeGatewayError without inspecting further. A nil
// return means the request is admissible to the session/server/registry
// admission path.
func websocketPreflight(request *http.Request) error {
	if err := validateWebSocketHandshake(request); err != nil {
		return err
	}
	if err := validateWebSocketOrigin(request); err != nil {
		return err
	}
	return nil
}

// writeWebSocketPreflightError maps a websocketPreflight return value to the
// pre-admission HTTP contract. A typed websocketPreflightError surfaces its
// own code/detail; any untyped error is normalised to
// INVALID_WEBSOCKET_HANDSHAKE so callers never observe an unstructured 500.
// WEBSOCKET_ORIGIN_REJECTED maps to 403; every other code maps to 400.
func writeWebSocketPreflightError(c *gin.Context, err error) {
	preflightError, ok := err.(websocketPreflightError)
	if !ok {
		preflightError = websocketPreflightError{
			code:   "INVALID_WEBSOCKET_HANDSHAKE",
			detail: err.Error(),
		}
	}
	switch preflightError.code {
	case "WEBSOCKET_ORIGIN_REJECTED":
		writeGatewayError(c, http.StatusForbidden, preflightError.code, preflightError.detail)
	default:
		writeGatewayError(c, http.StatusBadRequest, preflightError.code, preflightError.detail)
	}
}

// webSocketHandshakeHeaders are the request-side WS headers gorilla's Dialer
// rejects when explicitly set on the dialer's input header. cloneGatewayRequestHeaders
// already strips Connection and Upgrade; this list covers the Sec-Websocket-*
// family that gorilla also rejects. The downstream values must not bleed into
// the upstream dial.
var webSocketHandshakeHeaders = []string{
	"Sec-Websocket-Key",
	"Sec-Websocket-Version",
	"Sec-Websocket-Protocol",
	"Sec-Websocket-Extensions",
	"Sec-Websocket-Accept",
}

// stripWebSocketHandshakeHeaders removes the request-side WS handshake headers
// so the upstream dialer does not reject the sanitized header set.
func stripWebSocketHandshakeHeaders(header http.Header) {
	for _, key := range webSocketHandshakeHeaders {
		header.Del(key)
	}
}

// cloneWebSocketRequestHeader produces a sanitized header set suitable for the
// upstream WS dial: it inherits cloneGatewayRequestHeaders' hop-by-hop /
// sensitive-header stripping, then additionally removes the request-side WS
// headers gorilla's Dialer rejects. The returned header is a fresh copy that
// the caller fully owns.
func cloneWebSocketRequestHeader(source http.Header) http.Header {
	header := cloneGatewayRequestHeaders(source)
	stripWebSocketHandshakeHeaders(header)
	return header
}

// webSocketSubprotocolRequested reports whether selected appears in the
// requested list, compared case-insensitively per RFC 6455 §4.2.2.
func webSocketSubprotocolRequested(requested []string, selected string) bool {
	for _, candidate := range requested {
		if strings.EqualFold(candidate, selected) {
			return true
		}
	}
	return false
}

// proxyWebSocket runs the WebSocket gateway pipeline after preflight, session,
// server, registry, generation fence, and upstream validation have all passed.
//
// Phase 3 scope:
//   - canonical WS URL (http(s) → ws(s)) with the same query-rewrite rule the
//     HTTP path uses, sanitized headers stripped of WS-handshake + sensitive
//     fields, and client-requested subprotocols forwarded via gorilla's
//     Subprotocols list;
//   - 401 detection on the upstream dial with at most one refresh+retry;
//   - single-owner response.Body.Close invariant for every non-success path
//     and for the success-101 path after the pump exits;
//   - subprotocol verification: if the upstream's chosen Sec-WebSocket-Protocol
//     is outside the client's requested list, surface 502
//     WEBSOCKET_PROTOCOL_REJECTED before the TryCommit handoff;
//   - strict single-owner transition: upstream 101 + subprotocol verify →
//     entry.TryCommit → real downstream Hijack/Upgrade. A typed registry
//     cancellation that wins the TryCommit arbitration surfaces the typed
//     HTTP contract (409/503/no-write) without ever invoking the downstream
//     Upgrade;
//   - minimal block-on-downstream-read pump (Phase 4 replaces it with the
//     full bidirectional frame pump).
//
// The typed-cause stage matrix: after every blocking operation (UserToken,
// initial dial, RefreshUserToken, retry dial) the gateway runs
// writeTypedCancelOrGenerationChange — the unified pre-stage fence that
// catches BOTH a typed registry cause that has already landed on
// entry.Context() AND a session generation mismatch whose Cancel is
// in-flight from the watcher goroutine but has not yet reached the entry.
// The matrix is what makes TestGatewayWebSocketGenerationChangeDuringTokenFetch
// and TestGatewayWebSocketGenerationChangeDuringRefreshAbortsRetryDial
// deterministic: the hook's `select { ctx.Done / release }` race can no
// longer let a refresh-success path proceed to the retry dial.
func (gateway *desktopGateway) proxyWebSocket(
	c *gin.Context,
	query gatewayQuery,
	entry *streamregistry.Entry,
	requestContext context.Context,
	sessionID string,
	snapshot uisession.Snapshot,
	connection types.ServerGatewayConnection,
	upstream *url.URL,
	workspaceID string,
) {
	clientSubprotocols := websocket.Subprotocols(c.Request)

	userToken, err := connection.UserToken(requestContext)
	if err != nil {
		if gateway.writeTypedCancelOrGenerationChange(c, entry, sessionID, snapshot.Generation) {
			return
		}
		writeGatewayError(c, http.StatusBadGateway, "REMOTE_UNREACHABLE", "Unable to authenticate with the selected AgentTeam server.")
		return
	}
	if gateway.writeTypedCancelOrGenerationChange(c, entry, sessionID, snapshot.Generation) {
		return
	}

	upstreamConn, response, dialErr := gateway.dialUpstreamWebSocket(
		requestContext,
		c.Request,
		query,
		upstream,
		connection,
		userToken,
		clientSubprotocols,
		workspaceID,
	)

	if dialErr != nil && response != nil && response.StatusCode == http.StatusUnauthorized {
		// Single-owner body close for the 401 attempt; the refresh retry
		// owns its own response.Body lifecycle.
		_ = response.Body.Close()
		if gateway.writeTypedCancelOrGenerationChange(c, entry, sessionID, snapshot.Generation) {
			return
		}
		refreshedToken, refreshErr := connection.RefreshUserToken(requestContext, userToken)
		if refreshErr != nil {
			if gateway.writeTypedCancelOrGenerationChange(c, entry, sessionID, snapshot.Generation) {
				return
			}
			writeGatewayError(c, http.StatusBadGateway, "REMOTE_UNREACHABLE", "Unable to authenticate with the selected AgentTeam server.")
			return
		}
		if gateway.writeTypedCancelOrGenerationChange(c, entry, sessionID, snapshot.Generation) {
			return
		}
		userToken = refreshedToken
		upstreamConn, response, dialErr = gateway.dialUpstreamWebSocket(
			requestContext,
			c.Request,
			query,
			upstream,
			connection,
			userToken,
			clientSubprotocols,
			workspaceID,
		)
	}

	// Subprotocol mismatch: gorilla's Dialer rejects a server-chosen
	// protocol that was never offered by returning ErrBadHandshake with the
	// 101 response. Surface the typed 502 here so the downstream Upgrade is
	// never reached.
	if dialErr != nil && response != nil && response.StatusCode == http.StatusSwitchingProtocols {
		if selected := response.Header.Get("Sec-WebSocket-Protocol"); selected != "" &&
			!webSocketSubprotocolRequested(clientSubprotocols, selected) {
			_ = response.Body.Close()
			writeGatewayError(c, http.StatusBadGateway, "WEBSOCKET_PROTOCOL_REJECTED",
				"The upstream selected a subprotocol the client did not request.")
			return
		}
	}

	if dialErr != nil {
		if response != nil {
			_ = response.Body.Close()
		}
		if gateway.writeTypedCancelOrGenerationChange(c, entry, sessionID, snapshot.Generation) {
			return
		}
		writeGatewayError(c, http.StatusBadGateway, "REMOTE_UNREACHABLE", "Unable to reach the selected AgentTeam server.")
		return
	}

	// Post-dial fence (reviewer3 blocker #2): TryCommit only arbitrates
	// Cancel that has already landed on entry.Context(). The watcher-lag
	// window — generation changed but the proxy() watcher goroutine has
	// not yet called entry.Cancel — must be closed here so a generation
	// change (or typed cancel) that races with the just-returned dial
	// cannot silently proceed to subprotocol verify / TryCommit. Single
	// owner cleanup on hit: response.Body and upstreamConn each close
	// exactly once. This covers BOTH the initial and retry dial success
	// paths because they converge at this point.
	if gateway.writeTypedCancelOrGenerationChange(c, entry, sessionID, snapshot.Generation) {
		_ = response.Body.Close()
		_ = upstreamConn.Close()
		return
	}

	// Success: dial returned a live upstream conn. gorilla has already
	// verified the chosen subprotocol against clientSubprotocols, so the
	// extra check below is a defensive invariant against future regressions
	// in either gorilla or the local dial seam.
	selectedProtocol := upstreamConn.Subprotocol()
	if selectedProtocol != "" && !webSocketSubprotocolRequested(clientSubprotocols, selectedProtocol) {
		_ = response.Body.Close()
		_ = upstreamConn.Close()
		writeGatewayError(c, http.StatusBadGateway, "WEBSOCKET_PROTOCOL_REJECTED",
			"The upstream selected a subprotocol the client did not request.")
		return
	}

	// TryCommit: single owner-transfer from the registry-owned pre-commit
	// phase to the pump-owned post-commit phase. Arbitrated atomically
	// against Cancel. On (false, cause) the typed cause has already won —
	// close the upstream conn + Body (single owner) and surface the typed
	// HTTP contract WITHOUT invoking the downstream Upgrade commit step.
	// On (true, nil) the gateway owns the commit; the deferred closes
	// installed below take over single-owner duty for the pump lifetime.
	if owned, cause := entry.TryCommit(); !owned || cause != nil {
		if cause != nil {
			_ = response.Body.Close()
			_ = upstreamConn.Close()
			gateway.writePreCommitCancelCause(c, cause)
			return
		}
		// owned=false && cause=nil: entry was released before commit.
		// Close upstream resources and fail closed without inventing a
		// new HTTP contract.
		_ = response.Body.Close()
		_ = upstreamConn.Close()
		return
	}

	// Single owner of the upstream response.Body: deferred Close runs exactly
	// once regardless of which exit path the pump takes.
	defer func() {
		_ = response.Body.Close()
	}()
	defer func() {
		_ = upstreamConn.Close()
	}()

	// Upgrader is constructed AFTER TryCommit owns the commit step. Every
	// failure path from here on (Hijack rejection, header validation, etc.)
	// is commit-owned: the response body MUST NOT be rewritten with a
	// synthesized HTTP error. gorilla's Upgrader calls u.Error on Upgrade
	// failure; the default implementation http.Error's into c.Writer
	// (WriteHeader+Write), which would surface as a post-commit body
	// rewrite and break the Path C contract. Replace it with a no-op so
	// the only HTTP writes after TryCommit are the 101 switch-protocol
	// bytes that gorilla writes via the hijacked conn itself (which never
	// reach c.Writer's Write/Header methods).
	upgrader := websocket.Upgrader{
		CheckOrigin: func(*http.Request) bool { return true },
		Error:       func(http.ResponseWriter, *http.Request, int, error) {},
	}
	if selectedProtocol != "" {
		upgrader.Subprotocols = []string{selectedProtocol}
	}

	// Pre-upgrade cancel-watcher. Covers the window between TryCommit-
	// owned and the post-upgrade watcher taking over. The hijackCapturer
	// intercepts gorilla's internal Hijack() call so the watcher has the
	// raw downstream conn even while the 101 response Write is still
	// blocked. On entry.Context().Done() the watcher closes the raw conn
	// — this is what makes CancelAllInterruptsBlockedDownstream101
	// converge: the blocked Write fails immediately, gorilla's Upgrade
	// returns an error, the handler's deferred closes tear down the
	// upstream, and entry.Release runs.
	//
	// Path B/C zero HTTP rewrite invariant is preserved: the wrapper
	// delegates every ResponseWriter method, the no-op Error stays in
	// place, and the cancel path only touches the raw net.Conn (never
	// c.Writer). Observer does not change control flow: the watcher
	// exits when preUpgradeStop closes; the post-upgrade watcher then
	// owns the typed-close-frame job.
	capturer := &hijackCapturer{ResponseWriter: c.Writer}
	preUpgradeStop := make(chan struct{})
	go func() {
		select {
		case <-entry.Context().Done():
			capturer.closeRawConn()
		case <-preUpgradeStop:
		}
	}()

	downstreamConn, err := upgrader.Upgrade(capturer, c.Request, nil)
	close(preUpgradeStop)
	if err != nil {
		return
	}
	defer func() {
		_ = downstreamConn.Close()
	}()

	// Register the per-session lifecycle tracker AFTER Upgrade succeeds.
	// Every P5 worker (cancel-watcher, heartbeat worker, both pump
	// goroutines) spawns below this point, so the lifecycle signals are
	// guaranteed to all close on teardown. Registering before Upgrade
	// would expose half-initialized signals on the Upgrade-failure path
	// (closeWatcherClosed / heartbeatTickerClosed would never close
	// because their goroutines never spawned). The coordinator closes
	// its own signal via the LIFO-last defer so it fires AFTER the two
	// join-based defers below, which in turn wait for the worker
	// goroutines' stopped signals — statically proving the
	// "5 workers → coordinator → Entry.Release" ordering.
	lifecycle := newWSSessionLifecycle()
	gateway.registerLifecycle(sessionID, lifecycle)
	defer gateway.unregisterLifecycle(sessionID, lifecycle)
	defer close(lifecycle.coordinatorStopped)

	// Post-commit cancellation bridge: write the typed WS close frame
	// (1012 generation change, 1001 daemon shutdown / client disconnect
	// / deadline) to the downstream conn when entry.Context() fires, then
	// close the conn so the pump's NextReader unblocks and the handler
	// can run its deferred cleanup. WriteControl is concurrency-safe per
	// gorilla's documentation, so the cancel-watcher does not race the
	// pump's data-frame writes on the same conn.
	//
	// References are heap-allocated (*websocket.Conn and
	// *streamregistry.Entry), so the goroutine is race-safe against gin's
	// *gin.Context pool reset — unlike the HTTP path's deadline-interrupt
	// goroutine in proxy(), this watcher never touches c.Writer.
	//
	// Lifecycle: this is the close-watcher goroutine tracked by
	// LifecycleStopSignals.CloseWatcherStopped. The teardown defer
	// closes cancelWatcher AND joins on closeWatcherClosed so the
	// observer-invariant "coordinatorStopped fires after every worker
	// stopped signal" holds statically, not just empirically.
	cancelWatcher := make(chan struct{})
	go func() {
		defer close(lifecycle.closeWatcherClosed)
		select {
		case <-entry.Context().Done():
			cause := context.Cause(entry.Context())
			code, reason := gateway.closeCodeFromCause(cause)
			_ = downstreamConn.WriteControl(
				websocket.CloseMessage,
				websocket.FormatCloseMessage(code, reason),
				time.Now().Add(time.Second),
			)
			_ = downstreamConn.Close()
		case <-cancelWatcher:
		}
	}()
	defer func() {
		close(cancelWatcher)
		<-lifecycle.closeWatcherClosed
	}()

	// Heartbeat ticker worker. P5 wires ONLY the worker's lifecycle
	// exit so the observer's HeartbeatTickerStopped signal fires on
	// teardown; the ticker.C branch is a deliberate no-op so P5 does
	// not introduce new heartbeat wire semantics (WS PingMessage to
	// downstream) that have no contract or test coverage yet. A future
	// phase that adds real ping/pong health checks will populate the
	// tick branch with the wire-write behavior + dedicated test
	// coverage. Keeping the real *time.Ticker (instead of a bare
	// blocking goroutine) makes the HeartbeatTickerStopped signal
	// correspond to actual ticker resource lifecycle: teardown runs
	// ticker.Stop() before closing the stopped signal.
	heartbeatStop := make(chan struct{})
	go func() {
		defer close(lifecycle.heartbeatTickerClosed)
		ticker := time.NewTicker(webSocketHeartbeatInterval)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				// P5 no-op: real ping behavior lands in a later phase.
			case <-heartbeatStop:
				return
			}
		}
	}()
	defer func() {
		close(heartbeatStop)
		<-lifecycle.heartbeatTickerClosed
	}()

	// P4 bidirectional pump: per-direction single reader + single data
	// writer, within-message backpressure (source is not re-read until
	// the current chunk's Write completes; next message is not read
	// until the current writer is Closed), 1009 oversize termination
	// without FIN-complete truncation, and typed close-code mapping for
	// post-commit Cancel via the watcher above.
	//
	// Oversize is enforced authoritatively by copyWebSocketMessage's
	// cumulative limit check; gorilla SetReadLimit is intentionally NOT
	// installed so oversize is always surfaced through the copy path
	// where the no-Close/FIN contract is enforced.
	pump := &webSocketPump{
		gateway:    gateway,
		entry:      entry,
		downstream: downstreamConn,
		upstream:   upstreamConn,
		lifecycle:  lifecycle,
	}
	pump.run()
}

// dialUpstreamWebSocket builds the canonical WS URL (http(s) → ws(s)),
// applies the same query-rewrite rule as the HTTP path, sanitizes the
// downstream header set for gorilla's Dialer, and forwards the client's
// requested subprotocols. The returned (Conn, Response, error) tuple is
// the caller's responsibility: gorilla does not call Body.Close itself.
func (gateway *desktopGateway) dialUpstreamWebSocket(
	ctx context.Context,
	source *http.Request,
	query gatewayQuery,
	upstream *url.URL,
	connection types.ServerGatewayConnection,
	userToken string,
	subprotocols []string,
	workspaceID string,
) (*websocket.Conn, *http.Response, error) {
	target := *upstream
	target.Path = joinGatewayPath(upstream.Path, source.URL.Path)
	target.RawPath = joinGatewayPath(upstream.EscapedPath(), source.URL.EscapedPath())
	target.RawQuery = rewriteGatewayWebSocketQuery(query, userToken)
	switch target.Scheme {
	case "https":
		target.Scheme = "wss"
	default:
		target.Scheme = "ws"
	}
	header := cloneWebSocketRequestHeader(source.Header)
	header.Set("Authorization", "Bearer "+userToken)
	applyAuthoritativeWorkspaceHeadersForRequest(header, source, workspaceID)
	if workspaceID != "" && workspaceEventsStreamWorkspaceID(source.URL.Path) == workspaceID {
		header.Set("X-Workspace-Id", workspaceID)
		header.Set("X-AgentTeam-Route-Workspace-Id", workspaceID)
	}
	return connection.GatewayWebSocketDial(ctx, target.String(), header, subprotocols)
}

func workspaceEventsStreamWorkspaceID(path string) string {
	const prefix = "/api/v1/workspaces/"
	const suffix = "/events/stream"
	if !strings.HasPrefix(path, prefix) || !strings.HasSuffix(path, suffix) {
		return ""
	}
	workspaceID := strings.TrimSuffix(strings.TrimPrefix(path, prefix), suffix)
	if workspaceID == "" || strings.Contains(workspaceID, "/") {
		return ""
	}
	return workspaceID
}

// rewriteGatewayWebSocketQuery is the WS-specific credential injection. The
// HTTP path's query.rewrite is a no-op when credentialKey is empty (a
// deliberate HTTP-only decision to preserve query fidelity). The WS path
// must always carry the userToken so the upstream can authenticate the
// dial, so when credentialKey is empty we synthesize an accessToken entry
// while preserving any pre-existing query parameters.
func rewriteGatewayWebSocketQuery(query gatewayQuery, userToken string) string {
	if query.credentialKey != "" {
		return query.rewrite(userToken)
	}
	values := make(url.Values)
	for key, vs := range query.values {
		values[key] = append([]string(nil), vs...)
	}
	values.Set("accessToken", userToken)
	return values.Encode()
}