package api

import (
	"context"
	"errors"
	"io"
	"sync"
	"time"
	"unicode/utf8"

	"github.com/gorilla/websocket"

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

// webSocketMaxMessageBytes is the per-message size cap the P4 pump
// enforces in both directions. copyWebSocketMessage checks the
// cumulative total before every Write so an oversize payload never
// commits a FIN-complete truncated frame; gorilla SetReadLimit is
// intentionally NOT installed so oversize always surfaces through the
// copy path where the no-Close/FIN contract is enforced. The ErrReadLimit
// branch in pumpDirection is defensive — it keeps the pump correct if a
// future change reintroduces SetReadLimit.
const webSocketMaxMessageBytes int64 = 8 << 20

// webSocketHeartbeatInterval is the cadence at which the per-session
// heartbeat ticker fires WriteControl(PingMessage) on the downstream
// conn. P5 wires only the lifecycle exit so the observer's
// HeartbeatTickerStopped signal fires on teardown; real ping/pong
// health checks land in a later phase. The interval is long enough
// that the lifecycle test (which closes within ~milliseconds of
// connect) never observes a tick fire — observers only see the
// goroutine exit.
const webSocketHeartbeatInterval = 30 * time.Second

// errWebSocketMessageOversize is returned by copyWebSocketMessage when
// the cumulative byte count would exceed the per-message limit. The pump
// maps it — together with gorilla's ErrReadLimit — to a 1009 close frame
// on the downstream conn.
var errWebSocketMessageOversize = errors.New("websocket message exceeds size limit")

// copyWebSocketMessage is the per-message copy primitive shared by both
// pump directions. Production loops over gorilla NextReader/NextWriter
// and delegates one message at a time to this method, keeping the
// within-message backpressure invariant: the source is not read again
// until the current chunk's Write completes.
//
// Lazy-open semantics (planner3/codex5 P4 directive #2): the target
// writer is opened by openTarget AFTER the first non-empty Read clears
// the cumulative-limit check. This is what makes "first Read already
// exceeds the limit → target writer is never opened" hold statically.
// For an empty message (first Read returns 0 bytes + EOF) the writer
// IS opened and immediately Closed so the empty frame still commits.
//
// Cumulative-limit semantics (planner3 ruling, option A):
//   - Each non-empty Read is followed by the corresponding Write before
//     the next Read is issued.
//   - Before every Write, the cumulative total + chunk size is checked
//     against maxMessageBytes. If the chunk would exceed the limit, the
//     chunk is NOT written, the writer is NOT Opened or Closed (no FIN
//     commits a truncated frame), and errWebSocketMessageOversize is
//     returned.
//   - Already-written ≤ limit non-FIN prefix bytes are permitted; only
//     FIN-complete truncated messages are forbidden, and FIN is set
//     exclusively by Close on a clean EOF.
//
// Close discipline:
//   - Normal EOF (source Read returns io.EOF): if a writer has been
//     opened, target.Close is called exactly once so gorilla sets the
//     frame FIN bit. A nil error is returned.
//   - Source Read error other than EOF: the error is returned WITHOUT
//     calling Close. The caller tears down the connection; no FIN
//     commits a truncated frame.
//   - Target Write error: the error is returned WITHOUT calling Close.
//
// Gorilla's ErrReadLimit (defensive SetReadLimit layer) propagates
// unchanged; the pump maps it to a 1009 close frame.
func (gateway *desktopGateway) copyWebSocketMessage(
	source io.Reader,
	openTarget func() (io.WriteCloser, error),
	maxMessageBytes int64,
) error {
	if maxMessageBytes <= 0 {
		maxMessageBytes = webSocketMaxMessageBytes
	}
	buf := make([]byte, 32*1024)
	var total int64
	var target io.WriteCloser
	for {
		n, readErr := source.Read(buf)
		if n > 0 {
			if total+int64(n) > maxMessageBytes {
				return errWebSocketMessageOversize
			}
			if target == nil {
				opened, openErr := openTarget()
				if openErr != nil {
					return openErr
				}
				target = opened
			}
			wrote, writeErr := target.Write(buf[:n])
			if writeErr != nil {
				return writeErr
			}
			if wrote < n {
				return io.ErrShortWrite
			}
			total += int64(n)
		}
		if readErr != nil {
			if errors.Is(readErr, io.EOF) {
				if target == nil {
					opened, openErr := openTarget()
					if openErr != nil {
						return openErr
					}
					target = opened
				}
				if err := target.Close(); err != nil {
					return err
				}
				return nil
			}
			return readErr
		}
	}
}

// sanitizeWebSocketClose enforces RFC 6455 §5.5 close-frame constraints
// at the gateway boundary so an upstream or peer-supplied close payload
// never violates the wire contract:
//
//   - Status codes 1004, 1005, 1006, 1015 are reserved per RFC 6455
//     §7.4 and MUST NOT be sent on the wire. They are mapped to 1000
//     (normal closure) and the reason is cleared so a misleading
//     "hidden" reason does not reach the peer.
//   - The reason payload MUST be valid UTF-8 (RFC 6455 §5.5.1); invalid
//     UTF-8 is replaced with the empty string rather than forwarded.
//   - The close-frame payload cap is 123 bytes of reason (125-byte max
//     minus the 2-byte status code); an over-long reason is replaced
//     with the empty string so a misleading prefix does not leak.
//
// The returned (code, reason) is always safe to pass to gorilla's
// FormatCloseMessage.
func (gateway *desktopGateway) sanitizeWebSocketClose(code int, reason string) (int, string) {
	switch code {
	case websocket.CloseNoStatusReceived, // 1005
		websocket.CloseAbnormalClosure, // 1006
		1004,                           // RFC 6455 §7.4 reserved (no exported constant)
		websocket.CloseTLSHandshake:    // 1015
		code = websocket.CloseNormalClosure
		reason = ""
	}
	if !utf8.ValidString(reason) {
		reason = ""
	}
	if len(reason) > 123 {
		reason = ""
	}
	return code, reason
}

// closeCodeFromCause maps a streamregistry / context cancellation cause
// to the WebSocket close (code, reason) the post-commit watcher writes
// to the downstream conn. The mapping is part of the P4 typed-close
// contract:
//
//   - ReasonGenerationChange → 1012 ServiceRestart: the browser can
//     distinguish a server switch from a daemon shutdown and trigger
//     the correct reconnection flow.
//   - ReasonDaemonShutdown   → 1001 GoingAway: a clean daemon exit.
//   - ReasonClientDisconnect → 1001 GoingAway: the client is already
//     gone; the close frame is best-effort.
//   - context.DeadlineExceeded → 1001 GoingAway: the registry's
//     deadline bridge fired.
//
// An unrecognized cause falls back to 1001 GoingAway with the empty
// reason so the wire contract is never violated.
func (gateway *desktopGateway) closeCodeFromCause(cause error) (int, string) {
	switch {
	case errors.Is(cause, streamregistry.ReasonGenerationChange):
		return websocket.CloseServiceRestart, "server generation changed"
	case errors.Is(cause, streamregistry.ReasonDaemonShutdown):
		return websocket.CloseGoingAway, "daemon shutdown"
	case errors.Is(cause, streamregistry.ReasonClientDisconnect):
		return websocket.CloseGoingAway, "client disconnect"
	case errors.Is(cause, context.DeadlineExceeded):
		return websocket.CloseGoingAway, "deadline exceeded"
	}
	return websocket.CloseGoingAway, ""
}

// webSocketPump orchestrates the bidirectional message-level proxy
// between a downstream *websocket.Conn and an upstream *websocket.Conn.
// Per planner3 directive it is the single place that owns:
//
//   - One reader + one data-writer goroutine per direction.
//   - Within-message backpressure via copyWebSocketMessage: the source
//     is not read again until the current chunk's Write completes, and
//     the next message is not read until the current writer is Closed.
//   - Oversize termination: a 1009 close frame is sent to the
//     downstream conn; the offending target writer is NEVER Closed so
//     no FIN-complete truncated frame reaches the peer.
//   - Upstream close frame propagation: when source.NextReader returns
//     a *websocket.CloseError, the sanitized (code, reason) is
//     forwarded to the target so the peer observes the same close
//     semantics.
//
// The pump does NOT close the conns itself; the caller's deferred
// Close handles single-owner teardown. The pump's job is to exit both
// direction goroutines cleanly so the caller's defers run.
//
// Lifecycle: the downstream→upstream direction goroutine closes both
// downstreamReaderClosed and upstreamWriterClosed via defer when it
// exits — the same goroutine performs the synchronous Read→Write agentteam
// for that direction, so both roles stop together. The upstream→
// downstream direction is not individually tracked (it is supervised
// by the coordinator via the WaitGroup); the lifecycle observer only
// names the five signals in LifecycleStopSignals.
type webSocketPump struct {
	gateway    *desktopGateway
	entry      *streamregistry.Entry
	downstream *websocket.Conn
	upstream   *websocket.Conn
	lifecycle  *wsSessionLifecycle

	oversizeOnce sync.Once
}

// run blocks until both pump directions exit. The deferred teardown in
// proxyWebSocket runs after run returns.
func (p *webSocketPump) run() {
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		if p.lifecycle != nil {
			defer close(p.lifecycle.downstreamReaderClosed)
			defer close(p.lifecycle.upstreamWriterClosed)
		}
		p.pumpDirection(p.downstream, p.upstream)
	}()
	go func() {
		defer wg.Done()
		p.pumpDirection(p.upstream, p.downstream)
	}()
	wg.Wait()
}

// pumpDirection is the single-reader + single-data-writer agentteam for one
// direction. For each message:
//  1. source.NextReader returns the message type and a reader for its
//     payload.
//  2. The reader is handed to copyWebSocketMessage along with an
//     openTarget closure that calls target.NextWriter(messageType)
//     lazily — AFTER the first non-empty Read clears the cumulative
//     limit check. This is what makes the "first Read already exceeds
//     limit → target writer is never opened" half of the P4 contract
//     hold statically, not just empirically.
//  3. copyWebSocketMessage copies the bytes chunk-wise with cumulative
//     size limiting; Close is called on clean EOF, NEVER on error or
//     oversize.
//  4. The agentteam continues to the next message.
//
// On oversize (cumulative limit or gorilla ErrReadLimit) the pump fires
// a single 1009 close frame on the downstream conn via fireOversizeClose
// and exits. On any other error (source Read error, target Write error,
// close frame from peer) the agentteam exits; if the source returned a
// *websocket.CloseError, the sanitized (code, reason) is forwarded to
// the target so the peer observes matching close semantics.
//
// The deferred conn Close is what unblocks the peer direction's pending
// NextReader/NextWriter so the pump's run WaitGroup converges. Without
// it, an oversize on one direction would leave the peer direction
// parked on a NextReader that never returns (the peer will not send
// anything more after the oversize teardown), deadlocking the pump.
// gorilla's Close is idempotent (closeOnce), so the caller's deferred
// Close in proxyWebSocket remains a no-op.
func (p *webSocketPump) pumpDirection(source, target *websocket.Conn) {
	defer func() {
		_ = source.Close()
		_ = target.Close()
	}()
	for {
		messageType, reader, err := source.NextReader()
		if err != nil {
			if errors.Is(err, websocket.ErrReadLimit) {
				p.fireOversizeClose()
			} else {
				p.forwardCloseIfNeeded(target, err)
			}
			return
		}
		if messageType != websocket.TextMessage && messageType != websocket.BinaryMessage {
			// Gorilla surfaces Close/Ping/Pong control frames through
			// NextReader only when the application has not installed
			// handlers. Drain so the next data frame can be parsed; the
			// default Ping/Pong/Close handlers are sufficient for the
			// P4 contract.
			_, _ = copyAndDiscard(reader)
			continue
		}
		openTarget := func() (io.WriteCloser, error) {
			return target.NextWriter(messageType)
		}
		copyErr := p.gateway.copyWebSocketMessage(reader, openTarget, webSocketMaxMessageBytes)
		if copyErr == nil {
			continue
		}
		if errors.Is(copyErr, errWebSocketMessageOversize) || errors.Is(copyErr, websocket.ErrReadLimit) {
			p.fireOversizeClose()
		}
		return
	}
}

// forwardCloseIfNeeded inspects a NextReader error; if it carries a WS
// close frame (gorilla returns *websocket.CloseError on the receiving
// side), the sanitized (code, reason) is forwarded to target. Other
// errors are silent: the caller's deferred conn Close handles teardown.
//
// WriteControl is safe for concurrent use with message-frame writes per
// gorilla's documentation, so forwarding a close frame from one
// direction while the peer direction is mid-NextWriter does not race.
func (p *webSocketPump) forwardCloseIfNeeded(target *websocket.Conn, err error) {
	var closeErr *websocket.CloseError
	if !errors.As(err, &closeErr) {
		return
	}
	code, reason := p.gateway.sanitizeWebSocketClose(closeErr.Code, closeErr.Text)
	_ = target.WriteControl(
		websocket.CloseMessage,
		websocket.FormatCloseMessage(code, reason),
		time.Now().Add(time.Second),
	)
}

// fireOversizeClose writes a single 1009 close frame to the downstream
// conn. The sync.Once guarantees at most one 1009 is sent even when
// both directions race into oversize termination.
//
// WriteControl is concurrency-safe per gorilla's documentation.
func (p *webSocketPump) fireOversizeClose() {
	p.oversizeOnce.Do(func() {
		_ = p.downstream.WriteControl(
			websocket.CloseMessage,
			websocket.FormatCloseMessage(websocket.CloseMessageTooBig, ""),
			time.Now().Add(time.Second),
		)
	})
}

// copyAndDiscard drains r into io.Discard. It is a named helper so the
// pump's control-frame drain path reads as a single concept rather than
// an inlined io.Copy.
func copyAndDiscard(r io.Reader) (int64, error) {
	return io.Copy(io.Discard, r)
}