package websocket

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"github.com/pingcap-inc/nexuscrm-api-host/hostruntime"
	origincfg "github.com/pingcap-inc/nexuscrm-api-host/origin"
	"github.com/pingcap-inc/nexuscrm-contract/pkg/runtime/permissionstatev1"
	"github.com/rs/zerolog/log"
)

// wsClient wraps a WebSocket connection with write mutex
type wsClient struct {
	conn                   *websocket.Conn
	authorization          hostruntime.WebSocketAuthorizationLease
	authorizationTimeout   time.Duration
	writeTimeout           time.Duration
	authorizedUserID       string
	lastAuthorizationState hostruntime.AuthorizingPermissionState
	authorizationGate      chan struct{}
	retiring               atomic.Bool
	closed                 atomic.Bool
	writeMu                sync.Mutex
	closeOnce              sync.Once
	sessionCleanupOnce     sync.Once
}

func (c *wsClient) WriteJSON(v interface{}) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if c.writeTimeout > 0 {
		if err := c.conn.SetWriteDeadline(time.Now().Add(c.writeTimeout)); err != nil {
			return err
		}
	}
	return c.conn.WriteJSON(v)
}

func (c *wsClient) WriteMessage(messageType int, data []byte) error {
	c.writeMu.Lock()
	defer c.writeMu.Unlock()
	if c.writeTimeout > 0 {
		if err := c.conn.SetWriteDeadline(time.Now().Add(c.writeTimeout)); err != nil {
			return err
		}
	}
	return c.conn.WriteMessage(messageType, data)
}

func (c *wsClient) Close() {
	c.close(0, "")
}

func (c *wsClient) CloseWithMessage(code int, reason string) {
	c.close(code, reason)
}

func (c *wsClient) close(code int, reason string) {
	c.retiring.Store(true)
	c.closeOnce.Do(func() {
		c.closed.Store(true)
		if code != 0 {
			timeout := c.writeTimeout
			if timeout <= 0 || timeout > time.Second {
				timeout = time.Second
			}
			_ = c.conn.WriteControl(
				websocket.CloseMessage,
				websocket.FormatCloseMessage(code, reason),
				time.Now().Add(timeout),
			)
		}
		_ = c.conn.Close()
	})
}

// destroySessionWhenIdle keeps close non-blocking while ensuring a strong
// session is not destroyed underneath an already-admitted action.
func (c *wsClient) destroySessionWhenIdle(destroy func()) {
	c.sessionCleanupOnce.Do(func() {
		if c.authorization == nil {
			destroy()
			return
		}
		go func() {
			c.authorizationGate <- struct{}{}
			defer func() {
				<-c.authorizationGate
			}()
			destroy()
		}()
	})
}

func (c *wsClient) Authorize(
	ctx context.Context,
	request hostruntime.WebSocketAuthorizationRequest,
	action func() error,
) error {
	if action == nil {
		return &permissionstatev1.Failure{
			Category: permissionstatev1.FailureInvalidValue,
		}
	}
	if c.authorization == nil {
		return action()
	}
	if c.retiring.Load() {
		return errors.New("authorizing websocket: connection is closing")
	}
	if !request.Operation.Valid() ||
		request.Operation != hostruntime.WebSocketAuthorizeDisconnect &&
			strings.TrimSpace(request.Channel) == "" {
		return &permissionstatev1.Failure{
			Category: permissionstatev1.FailureInvalidValue,
		}
	}
	if ctx == nil {
		ctx = context.Background()
	}
	ctx, cancel := context.WithTimeout(ctx, c.authorizationTimeout)
	defer cancel()

	select {
	case c.authorizationGate <- struct{}{}:
		defer func() {
			<-c.authorizationGate
		}()
	case <-ctx.Done():
		c.retiring.Store(true)
		return ctx.Err()
	}
	if c.retiring.Load() || c.closed.Load() {
		return errors.New("authorizing websocket: connection is closing")
	}
	state, err := reauthorizeWebSocket(ctx, c.authorization, request)
	if err != nil {
		c.retiring.Store(true)
		return err
	}
	if _, err := hostruntime.ContextWithAuthorizingPermissionState(ctx, state); err != nil {
		c.retiring.Store(true)
		return err
	}
	if state.Subject.UserID() != c.authorizedUserID {
		c.retiring.Store(true)
		return &permissionstatev1.Failure{
			Category: permissionstatev1.FailureInvalidValue,
		}
	}
	current := state.Snapshot.Version()
	previous := c.lastAuthorizationState.Snapshot.Version()
	comparison, err := current.Compare(previous)
	if err != nil {
		c.retiring.Store(true)
		return err
	}
	if comparison < 0 {
		c.retiring.Store(true)
		return &permissionstatev1.Failure{
			Category: permissionstatev1.FailureStaleSnapshot,
		}
	}
	if comparison == 0 &&
		(!slices.Equal(
			state.Snapshot.SystemPermissions(),
			c.lastAuthorizationState.Snapshot.SystemPermissions(),
		) ||
			!slices.Equal(
				state.Snapshot.CustomPermissions(),
				c.lastAuthorizationState.Snapshot.CustomPermissions(),
			)) {
		c.retiring.Store(true)
		return &permissionstatev1.Failure{
			Category: permissionstatev1.FailureAuthorityUnavailable,
		}
	}
	if err := ctx.Err(); err != nil {
		c.retiring.Store(true)
		return err
	}
	c.lastAuthorizationState = state
	if c.retiring.Load() || c.closed.Load() {
		return errors.New("authorizing websocket: connection is closing")
	}
	if err := action(); err != nil {
		c.retiring.Store(true)
		return err
	}
	return nil
}

type webSocketReauthorizationResult struct {
	state hostruntime.AuthorizingPermissionState
	err   error
}

func reauthorizeWebSocket(
	ctx context.Context,
	authorization hostruntime.WebSocketAuthorizationLease,
	request hostruntime.WebSocketAuthorizationRequest,
) (hostruntime.AuthorizingPermissionState, error) {
	result := make(chan webSocketReauthorizationResult, 1)
	go func() {
		state, err := authorization.ReauthorizeWebSocket(ctx, request)
		result <- webSocketReauthorizationResult{state: state, err: err}
	}()
	select {
	case completed := <-result:
		if err := ctx.Err(); err != nil {
			return hostruntime.AuthorizingPermissionState{}, err
		}
		return completed.state, completed.err
	case <-ctx.Done():
		return hostruntime.AuthorizingPermissionState{}, ctx.Err()
	}
}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     checkWebSocketOrigin,
}

// checkWebSocketOrigin validates the origin for WebSocket connections
func checkWebSocketOrigin(r *http.Request) bool {
	origin := r.Header.Get("Origin")
	if origin == "" {
		return true // Allow requests without Origin header (e.g., curl)
	}

	if origincfg.IsAllowedOrigin(origin) {
		return true
	}

	log.Warn().Str("origin", origin).Msg("WebSocket connection rejected: origin not allowed")
	return false
}

// WebSocketHandler handles WebSocket connections for real-time streaming
type WebSocketHandler struct {
	sessionManager         hostruntime.WebSocketSessionStore
	featureGate            hostruntime.FeatureGate
	authorizationFactory   hostruntime.WebSocketAuthorizationLeaseFactory
	reauthorizationTimeout time.Duration
	writeTimeout           time.Duration
	clients                map[string]*wsClient
	mu                     sync.RWMutex
}

// NewWebSocketHandler creates a new WebSocket handler
func NewWebSocketHandler(sessionMgr hostruntime.WebSocketSessionStore, featureGate hostruntime.FeatureGate) *WebSocketHandler {
	return &WebSocketHandler{
		sessionManager: sessionMgr,
		featureGate:    featureGate,
		clients:        make(map[string]*wsClient),
	}
}

// AuthorizingWebSocketHandlerOptions closes the long-lived authorization
// composition. ReauthorizationTimeout is an operational upper bound for each
// subscribe, unsubscribe, disconnect or delivery decision. WriteTimeout bounds
// every protected response, delivery and control write.
type AuthorizingWebSocketHandlerOptions struct {
	SessionManager         hostruntime.WebSocketSessionStore
	FeatureGate            hostruntime.FeatureGate
	AuthorizationFactory   hostruntime.WebSocketAuthorizationLeaseFactory
	ReauthorizationTimeout time.Duration
	WriteTimeout           time.Duration
}

// NewAuthorizingWebSocketHandler constructs a WebSocket handler that rejects
// incomplete strong composition. The legacy constructor remains unchanged for
// source compatibility but cannot activate this path.
func NewAuthorizingWebSocketHandler(
	options AuthorizingWebSocketHandlerOptions,
) (*WebSocketHandler, error) {
	if options.SessionManager == nil {
		return nil, errors.New(
			"authorizing websocket: session manager is required",
		)
	}
	if options.AuthorizationFactory == nil {
		return nil, errors.New(
			"authorizing websocket: authorization lease factory is required",
		)
	}
	if options.ReauthorizationTimeout <= 0 {
		return nil, errors.New(
			"authorizing websocket: positive reauthorization timeout is required",
		)
	}
	if options.WriteTimeout <= 0 {
		return nil, errors.New(
			"authorizing websocket: positive write timeout is required",
		)
	}
	return &WebSocketHandler{
		sessionManager:         options.SessionManager,
		featureGate:            options.FeatureGate,
		authorizationFactory:   options.AuthorizationFactory,
		reauthorizationTimeout: options.ReauthorizationTimeout,
		writeTimeout:           options.WriteTimeout,
		clients:                make(map[string]*wsClient),
	}, nil
}

// HandleConnection upgrades HTTP to WebSocket and manages the connection
// @Summary WebSocket connection endpoint
// @Tags Streaming
// @Success 101 {string} string "Switching Protocols"
// @Router /ws [get]
func (h *WebSocketHandler) HandleConnection(c *gin.Context) {
	if h.featureGate != nil && h.featureGate.Disabled(hostruntime.FeatureWebSocket) {
		log.Warn().Msg("WebSocket feature is disabled")
		c.Status(http.StatusServiceUnavailable)
		return
	}
	if h.sessionManager == nil {
		log.Error().Msg("WebSocket session manager is not configured")
		c.Status(http.StatusServiceUnavailable)
		return
	}
	authorization, initialState, err := h.openAuthorization(c)
	if err != nil {
		status := http.StatusServiceUnavailable
		if errors.Is(err, hostruntime.ErrAuthorizingCredentialRejected) {
			status = http.StatusUnauthorized
		}
		log.Error().Int("status", status).Msg(
			"WebSocket authorization lease could not be established",
		)
		c.Status(status)
		return
	}

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		log.Error().Err(err).Msg("WebSocket upgrade failed")
		return
	}

	// Create session for this WebSocket connection
	clientID := h.sessionManager.CreateSession()
	client := &wsClient{
		conn:                   conn,
		authorization:          authorization,
		authorizationTimeout:   h.reauthorizationTimeout,
		writeTimeout:           h.writeTimeout,
		authorizedUserID:       initialState.Subject.UserID(),
		lastAuthorizationState: initialState,
		authorizationGate:      make(chan struct{}, 1),
	}

	h.mu.Lock()
	h.clients[clientID] = client
	h.mu.Unlock()

	defer h.removeClient(clientID, client)

	log.Info().
		Str("client_id", clientID).
		Str("remote_addr", c.ClientIP()).
		Msg("WebSocket client connected")

	// Send handshake response
	handshakeMsg := map[string]interface{}{
		"channel":    "/meta/handshake",
		"clientId":   clientID,
		"successful": true,
		"version":    "1.0",
	}
	if err := client.WriteJSON(handshakeMsg); err != nil {
		log.Error().Err(err).Msg("Failed to send handshake")
		return
	}

	// Set read deadline for ping/pong
	conn.SetReadDeadline(time.Now().Add(120 * time.Second))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(120 * time.Second))
		return nil
	})

	// Start ping ticker
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	done := make(chan struct{})

	// Handle incoming messages
	go func() {
		defer close(done)
		for {
			_, message, err := conn.ReadMessage()
			if err != nil {
				if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
					log.Warn().Err(err).Str("client_id", clientID).Msg("WebSocket read error")
				}
				return
			}

			if err := h.handleMessage(
				c.Request.Context(),
				clientID,
				client,
				message,
			); err != nil {
				h.closeForAuthorizationFailure(clientID, client, err)
				return
			}
		}
	}()

	// Send pings
	for {
		select {
		case <-done:
			return
		case <-ticker.C:
			if err := client.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}
		}
	}
}

func (h *WebSocketHandler) openAuthorization(
	c *gin.Context,
) (
	hostruntime.WebSocketAuthorizationLease,
	hostruntime.AuthorizingPermissionState,
	error,
) {
	if h.authorizationFactory == nil {
		return nil, hostruntime.AuthorizingPermissionState{}, nil
	}
	state, ok := hostruntime.AuthorizingPermissionStateFromContext(
		c.Request.Context(),
	)
	if !ok {
		return nil, hostruntime.AuthorizingPermissionState{}, errors.New(
			"authorizing websocket: initial Permission State is required",
		)
	}
	if _, err := hostruntime.ContextWithAuthorizingPermissionState(
		c.Request.Context(),
		state,
	); err != nil {
		return nil, hostruntime.AuthorizingPermissionState{}, fmt.Errorf(
			"authorizing websocket: invalid initial Permission State: %w",
			err,
		)
	}
	bearer, ok := webSocketBearerToken(
		c.GetHeader("Authorization"),
	)
	if !ok {
		return nil, hostruntime.AuthorizingPermissionState{}, errors.New(
			"authorizing websocket: admitted bearer is required",
		)
	}
	openContext, cancel := context.WithTimeout(
		c.Request.Context(),
		h.reauthorizationTimeout,
	)
	defer cancel()
	authorization, err := openWebSocketAuthorization(
		openContext,
		h.authorizationFactory,
		bearer,
		state,
	)
	if err != nil {
		return nil, hostruntime.AuthorizingPermissionState{}, err
	}
	if authorization == nil {
		return nil, hostruntime.AuthorizingPermissionState{}, errors.New(
			"authorizing websocket: authorization lease is required",
		)
	}
	c.Request.Header.Del("Authorization")
	return authorization, state, nil
}

type webSocketAuthorizationOpenResult struct {
	authorization hostruntime.WebSocketAuthorizationLease
	err           error
}

func openWebSocketAuthorization(
	ctx context.Context,
	factory hostruntime.WebSocketAuthorizationLeaseFactory,
	bearer string,
	initial hostruntime.AuthorizingPermissionState,
) (hostruntime.WebSocketAuthorizationLease, error) {
	result := make(chan webSocketAuthorizationOpenResult, 1)
	go func() {
		authorization, err := factory.OpenWebSocketAuthorization(
			ctx,
			bearer,
			initial,
		)
		result <- webSocketAuthorizationOpenResult{
			authorization: authorization,
			err:           err,
		}
	}()
	select {
	case completed := <-result:
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		return completed.authorization, completed.err
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func webSocketBearerToken(header string) (string, bool) {
	parts := strings.SplitN(header, " ", 2)
	if len(parts) != 2 ||
		parts[0] != "Bearer" ||
		strings.TrimSpace(parts[1]) == "" {
		return "", false
	}
	return parts[1], true
}

// handleMessage processes incoming WebSocket messages
func (h *WebSocketHandler) handleMessage(
	ctx context.Context,
	clientID string,
	client *wsClient,
	message []byte,
) error {
	var msg map[string]interface{}
	if err := json.Unmarshal(message, &msg); err != nil {
		log.Warn().Err(err).Msg("Invalid WebSocket message format")
		return nil
	}

	channel, _ := msg["channel"].(string)

	switch channel {
	case "/meta/subscribe":
		subscription, _ := msg["subscription"].(string)
		if err := client.Authorize(
			ctx,
			hostruntime.WebSocketAuthorizationRequest{
				Operation: hostruntime.WebSocketAuthorizeSubscribe,
				Channel:   subscription,
			},
			func() error {
				h.sessionManager.Subscribe(clientID, subscription)
				return client.WriteJSON(map[string]interface{}{
					"channel":      "/meta/subscribe",
					"subscription": subscription,
					"successful":   true,
				})
			},
		); err != nil {
			return err
		}
		log.Debug().Str("client_id", clientID).Str("channel", subscription).Msg("WebSocket subscription added")

	case "/meta/unsubscribe":
		subscription, _ := msg["subscription"].(string)
		if err := client.Authorize(
			ctx,
			hostruntime.WebSocketAuthorizationRequest{
				Operation: hostruntime.WebSocketAuthorizeUnsubscribe,
				Channel:   subscription,
			},
			func() error {
				h.sessionManager.Unsubscribe(clientID, subscription)
				return client.WriteJSON(map[string]interface{}{
					"channel":      "/meta/unsubscribe",
					"subscription": subscription,
					"successful":   true,
				})
			},
		); err != nil {
			return err
		}

	case "/meta/disconnect":
		if err := client.Authorize(
			ctx,
			hostruntime.WebSocketAuthorizationRequest{
				Operation: hostruntime.WebSocketAuthorizeDisconnect,
			},
			func() error {
				return client.WriteJSON(map[string]interface{}{
					"channel":    "/meta/disconnect",
					"successful": true,
				})
			},
		); err != nil {
			return err
		}
	}
	return nil
}

func (h *WebSocketHandler) closeForAuthorizationFailure(
	clientID string,
	client *wsClient,
	err error,
) {
	code := websocket.CloseInternalServerErr
	reason := "authorization unavailable"
	if errors.Is(err, hostruntime.ErrAuthorizingCredentialRejected) {
		code = websocket.ClosePolicyViolation
		reason = "authorization rejected"
	}
	client.CloseWithMessage(code, reason)
	h.removeClient(clientID, client)
}

func (h *WebSocketHandler) removeClient(
	clientID string,
	client *wsClient,
) {
	h.mu.Lock()
	current, exists := h.clients[clientID]
	if exists && current == client {
		delete(h.clients, clientID)
	}
	h.mu.Unlock()
	if !exists || current != client {
		return
	}
	client.Close()
	client.destroySessionWhenIdle(func() {
		h.sessionManager.DestroySession(clientID)
	})
}

// Broadcast sends a message to all clients subscribed to a channel
func (h *WebSocketHandler) Broadcast(channel string, data interface{}) {
	h.mu.RLock()
	clientsCopy := make(map[string]*wsClient, len(h.clients))
	for id, c := range h.clients {
		clientsCopy[id] = c
	}
	h.mu.RUnlock()

	message := map[string]interface{}{
		"channel": channel,
		"data":    data,
	}

	for clientID, client := range clientsCopy {
		// Check if client is subscribed to this channel
		channels := h.sessionManager.GetSubscribedChannels(clientID)
		for _, ch := range channels {
			if ch == channel {
				if err := client.Authorize(
					context.Background(),
					hostruntime.WebSocketAuthorizationRequest{
						Operation: hostruntime.WebSocketAuthorizeDelivery,
						Channel:   channel,
					},
					func() error {
						return client.WriteJSON(message)
					},
				); err != nil {
					h.closeForAuthorizationFailure(
						clientID,
						client,
						err,
					)
					break
				}
				break
			}
		}
	}
}

// ClientCount returns the number of connected WebSocket clients
func (h *WebSocketHandler) ClientCount() int {
	h.mu.RLock()
	defer h.mu.RUnlock()
	return len(h.clients)
}