package websocket

import (
	"context"
	"fmt"
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/go-admin-team/go-admin-core/logger"
	"github.com/go-admin-team/go-admin-core/sdk/pkg/jwtauth"
	"github.com/gorilla/websocket"
)

var upgrader = websocket.Upgrader{
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许跨域
	},
}

func ServeWS(hub *Hub) gin.HandlerFunc {
	return func(c *gin.Context) {
		// 这里假设 JWT 中间件已验证，且用户 ID 在 Context
		claims := jwtauth.ExtractClaims(c)
		userID, ok := claims["identity"].(float64)
		if !ok {
			c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid user ID in token"})
			return
		}

		// 设置 Subprotocols 支持客户端传来的协议（token等）
		upgrader.Subprotocols = []string{c.GetHeader("Sec-WebSocket-Protocol")}

		conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			logger.Errorf("WebSocket upgrade failed: %v", err)
			c.JSON(http.StatusBadRequest, gin.H{"error": "upgrade failed"})
			return
		}

		ctx, cancel := context.WithCancel(context.Background())

		client := &Client{
			ID:         fmt.Sprintf("%v", userID),
			Conn:       conn,
			Send:       make(chan []byte, 1024),
			Context:    ctx,
			CancelFunc: cancel,
		}

		hub.Register <- client

		go client.Read(hub)
		go client.Write()
	}
}