package middleware_ws

import (
	"errors"
	"gitee.com/os-lee/easy-paas/common"
	"gitee.com/os-lee/easy-paas/common/e"
	"gitee.com/os-lee/easy-paas/gateway/internal/dao"
	"gitee.com/os-lee/easy-paas/gateway/internal/load_balancer"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"net/http"
	"net/url"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		// 允许任何来源的请求
		return true
	},
}

// Proxy 处理代理的WebSocket连接
func Proxy() gin.HandlerFunc {
	return func(c *gin.Context) {
		// 检查请求头是否为WebSocket请求
		if c.Request.Header.Get("Upgrade") != "websocket" {
			c.AbortWithStatus(http.StatusBadRequest)
			return
		}

		// 解析目标URL
		serverInterface, ok := c.Get("service")
		if !ok {
			common.ResponseError(c, e.ErrGateway, errors.New("从上下文中获取服务失败"))
			c.AbortWithStatus(http.StatusBadRequest)
			return
		}
		serviceDetail := serverInterface.(*dao.ServiceDetail)
		lb, err := load_balancer.Handler.Get(serviceDetail)
		targetURL, err := url.Parse(lb.Pick(c.Request).URL.String())
		if err != nil || targetURL.Scheme == "" || targetURL.Host == "" {
			c.AbortWithStatus(http.StatusBadRequest)
			return
		}

		// 升级HTTP连接到WebSocket
		src, err := upgrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			c.AbortWithError(http.StatusInternalServerError, err)
			return
		}
		defer src.Close()

		// 连接到远程服务器
		dialer := &websocket.Dialer{}
		dest, _, err := dialer.Dial(targetURL.String(), nil)
		if err != nil {
			c.AbortWithError(http.StatusInternalServerError, err)
			return
		}
		defer dest.Close()

		forwardMessages(src, dest)
		c.Abort()
	}
}

// 在两个WebSocket连接之间转发消息
func forwardMessages(src *websocket.Conn, dest *websocket.Conn) {
	srcReadChan := make(chan []byte, 10)
	destReadChan := make(chan []byte, 10)

	go readMessages(src, srcReadChan)
	go readMessages(dest, destReadChan)

	for {
		select {
		case data := <-srcReadChan:
			if err := dest.WriteMessage(websocket.TextMessage, data); err != nil {
				break
			}
		case data := <-destReadChan:
			if err := src.WriteMessage(websocket.TextMessage, data); err != nil {
				break
			}
		}
	}
}

// 从WebSocket连接读取消息并放入channel
func readMessages(conn *websocket.Conn, ch chan<- []byte) {
	for {
		_, data, err := conn.ReadMessage()
		if err != nil {
			close(ch)
			return
		}
		ch <- data
	}
}