package localruntime

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"strings"
	"time"

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

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

// activeStreamAuthRequiredClose is the application-level close code the
// server sends after upgrading when the session token is missing or
// invalid. The 4xxx range is reserved for application-defined close
// codes per RFC 6455 §7.4.2; 4401 is unambiguous with HTTP 401 and
// never collides with a protocol-level code. Browsers cannot read the
// HTTP handshake status of a failed WS upgrade, so we always upgrade
// and then surface auth failure via this close code so the client can
// distinguish it from a transport drop.
const activeStreamAuthRequiredClose = 4401

// activeStreamPongWait is how long the server waits between pong frames
// from the client before treating the connection as dead. The client is
// expected to respond to PingMessage with PongMessage within this window.
const activeStreamPongWait = 60 * time.Second

// activeStreamPingPeriod is the cadence at which the server sends
// PingMessage control frames. It must be shorter than activeStreamPongWait
// so a single dropped ping still leaves room for the pong deadline.
const activeStreamPingPeriod = 30 * time.Second

// activeStreamWriteWait is the per-write deadline used for both data
// frames and control frames so a stuck client cannot park a goroutine
// forever.
const activeStreamWriteWait = 10 * time.Second

// activeStreamHelloTone bounds how long the WS handshake may take before
// gorilla rolls it back. It also serves as HandshakeTimeout.
const activeStreamHelloTone = 5 * time.Second

var activeStreamDataWriteWait = activeStreamWriteWait

var activeStreamUpgrader = websocket.Upgrader{
	CheckOrigin: func(request *http.Request) bool {
		return validateActiveStreamOrigin(request)
	},
	Error: func(w http.ResponseWriter, _ *http.Request, status int, _ error) {
		// Keep the response body generic: the request URL may contain a session
		// token and must never be echoed into an error or log.
		http.Error(w, http.StatusText(status), status)
	},
	HandshakeTimeout: activeStreamHelloTone,
}

func validateActiveStreamOrigin(request *http.Request) bool {
	if request == nil {
		return false
	}
	return localsecurity.ValidateWebSocketOrigin(request) == nil
}

// activeSelectionStreamFrame is the on-wire frame shape. Type identifies
// the payload kind (hello|snapshot|selected|cleared|advanced); revision
// is the uisession revision at emission time so the client can deduplicate
// after reconnect. Server/Binding are populated only for "selected"-state
// frames.
type activeSelectionStreamFrame struct {
	Type       string                       `json:"type"`
	Revision   uint64                       `json:"revision,omitempty"`
	State      string                       `json:"state,omitempty"`
	Generation uint64                       `json:"generation,omitempty"`
	Server     *activeServerPublicResponse  `json:"server,omitempty"`
	Binding    *activeServerBindingResponse `json:"binding,omitempty"`
}

// StreamHandlerDeps bundles the WS handler's dependencies.
type StreamHandlerDeps struct {
	Provider SubscriptionProvider
	Servers  ServerRegistry
	Timer    LeaseTimerRegistry
}

// NewActiveSelectionStreamHandler upgrades the request at
// /servers/active/stream to a WebSocket and streams active-server
// selection events. The wire contract is JSON-over-text-frames so the
// desktop webview keeps one long-lived connection open for the lifetime
// of the window instead of re-establishing an EventSource per
// navigation.
//
// Auth: browsers cannot read the HTTP handshake status of a failed WS
// upgrade, so auth is enforced AFTER upgrade. The handshake MUST carry
// a valid session token via the Authorization Bearer header,
// X-Session-Token header, or accessToken query parameter. A missing or
// invalid token surfaces as a 4401 close frame
// (activeStreamAuthRequiredClose) so the client's onclose fires
// synchronously and the reconnect agentteam maps it to onTerminalAuth (full
// re-bootstrap) instead of retrying.
//
// Concurrency: all writes to the WS conn happen on the handler goroutine
// — data frames via WriteMessage, control frames via WriteControl. A
// single reader goroutine drains client frames to refresh the pong
// deadline; its exit (Close, transport error) signals the writer agentteam
// via readerDone.
func NewActiveSelectionStreamHandler(deps StreamHandlerDeps) gin.HandlerFunc {
	return func(c *gin.Context) {
		if deps.Provider == nil || deps.Servers == nil {
			internalUnavailable(c)
			return
		}

		conn, err := activeStreamUpgrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			// Upgrade already wrote an HTTP error response; nothing more to do.
			return
		}
		defer conn.Close()
		conn.SetReadDeadline(time.Now().Add(activeStreamPongWait))
		conn.SetPongHandler(func(string) error {
			return conn.SetReadDeadline(time.Now().Add(activeStreamPongWait))
		})

		if !authenticateActiveSelectionStream(c) {
			writeActiveStreamClose(conn, activeStreamAuthRequiredClose, "auth required")
			return
		}
		sessionID, expiresAt, ok := auth.SessionIdentity(c)
		if !ok {
			writeActiveStreamClose(conn, activeStreamAuthRequiredClose, "auth required")
			return
		}

		subscription, err := deps.Provider.Subscribe(sessionID, expiresAt)
		if err != nil {
			switch {
			case errors.Is(err, uisession.ErrSessionExpired),
				errors.Is(err, uisession.ErrSessionClosed),
				errors.Is(err, uisession.ErrSessionRequired):
				writeActiveStreamClose(conn, activeStreamAuthRequiredClose, "session required")
			default:
				writeActiveStreamClose(conn, websocket.CloseGoingAway, "runtime unavailable")
			}
			return
		}
		release := subscription.Release
		if release == nil {
			release = func() {}
		}
		defer release()

		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()
		stopContextRelease := context.AfterFunc(ctx, release)
		defer stopContextRelease()

		timers := deps.Timer
		if timers == nil {
			timers = realLeaseTimerRegistry{}
		}
		stopLease := timers.AfterFunc(time.Until(expiresAt), func() {
			release()
			cancel()
		})
		defer stopLease()

		if !sendActiveSelectionFrame(conn, activeSelectionStreamFrame{Type: "hello"}) {
			return
		}
		snapshotFrame, ok := buildActiveSelectionFrame("snapshot", subscription.Snapshot.State,
			subscription.Snapshot.ActiveServerID, subscription.Snapshot.Generation,
			subscription.Snapshot.Revision, deps.Servers)
		if !ok {
			return
		}
		if !sendActiveSelectionFrame(conn, snapshotFrame) {
			return
		}

		pingTicker := time.NewTicker(activeStreamPingPeriod)
		defer pingTicker.Stop()

		readerDone := make(chan struct{})
		go func() {
			defer close(readerDone)
			for {
				if _, _, err := conn.NextReader(); err != nil {
					return
				}
			}
		}()

		for {
			select {
			case <-ctx.Done():
				return
			case <-readerDone:
				return
			case <-pingTicker.C:
				if !writeActiveStreamPing(conn) {
					return
				}
			case change, open := <-subscription.Changes:
				if !open {
					writeActiveStreamClose(conn, websocket.CloseGoingAway, "session closed")
					return
				}
				kind, ok := mapChangeKind(change.Kind)
				if !ok {
					return
				}
				frame, ok := buildActiveSelectionFrame(kind, change.State, change.ActiveServerID,
					change.Generation, change.Revision, deps.Servers)
				if !ok {
					return
				}
				if !sendActiveSelectionFrame(conn, frame) {
					return
				}
			}
		}
	}
}

// authenticateActiveSelectionStream accepts the accessToken query
// parameter (browsers cannot set headers on a WS handshake) plus the
// Authorization/X-Session-Token headers used by non-browser clients.
// A token in any of the three locations is sufficient.
func authenticateActiveSelectionStream(c *gin.Context) bool {
	if _, _, ok := auth.SessionIdentity(c); ok {
		return true
	}
	if err := localsecurity.ValidateQueryCredentials(c.Request); err != nil {
		return false
	}
	tokenValue := strings.TrimSpace(c.Query("accessToken"))
	if tokenValue == "" {
		if header := c.GetHeader("Authorization"); header != "" {
			tokenValue = strings.TrimPrefix(header, "Bearer ")
		}
	}
	if tokenValue == "" {
		tokenValue = c.GetHeader("X-Session-Token")
	}
	if tokenValue == "" {
		return false
	}
	return auth.AuthenticateSession(c, tokenValue)
}

func buildActiveSelectionFrame(
	kind string,
	state uisession.SubscriptionState,
	serverID string,
	generation uint64,
	revision uint64,
	servers ServerRegistry,
) (activeSelectionStreamFrame, bool) {
	stateName, ok := mapSubscriptionState(state)
	if !ok {
		return activeSelectionStreamFrame{}, false
	}
	frame := activeSelectionStreamFrame{
		Type:       kind,
		Revision:   revision,
		State:      stateName,
		Generation: generation,
	}
	if state != uisession.SubscriptionStateSelected {
		return frame, true
	}
	server, ok := resolveStreamServer(servers, serverID)
	if !ok {
		return activeSelectionStreamFrame{}, false
	}
	frame.Server = &activeServerPublicResponse{
		ID:          server.ID,
		Name:        server.Name,
		Description: server.Description,
	}
	frame.Binding = &activeServerBindingResponse{
		WorkspaceID: server.WorkspaceID,
		MachineID:   server.MachineID,
	}
	return frame, true
}

func sendActiveSelectionFrame(conn *websocket.Conn, frame activeSelectionStreamFrame) bool {
	body, err := json.Marshal(frame)
	if err != nil {
		return false
	}
	if err := conn.SetWriteDeadline(time.Now().Add(activeStreamDataWriteWait)); err != nil {
		return false
	}
	if err := conn.WriteMessage(websocket.TextMessage, body); err != nil {
		return false
	}
	return true
}

func writeActiveStreamPing(conn *websocket.Conn) bool {
	if err := conn.WriteControl(
		websocket.PingMessage,
		nil,
		time.Now().Add(activeStreamWriteWait),
	); err != nil {
		return false
	}
	return true
}

func writeActiveStreamClose(conn *websocket.Conn, code int, reason string) {
	_ = conn.WriteControl(
		websocket.CloseMessage,
		websocket.FormatCloseMessage(code, reason),
		time.Now().Add(activeStreamWriteWait),
	)
}