package websocket

import (
	"encoding/json"
	"time"

	"github.com/flipped-aurora/gin-vue-admin/server/global"
	"github.com/gin-gonic/gin"
	"github.com/olahol/melody"
)

// HandleConnection 处理WebSocket连接
func HandleConnection(ctx *gin.Context) {
	// 从中间件获取用户ID
	userId, exists := ctx.Get("userId")
	if !exists {
		ctx.JSON(401, gin.H{"error": "用户认证失败"})
		return
	}

	// 升级为WebSocket连接，并存储用户信息
	global.MELODY.HandleRequestWithKeys(ctx.Writer, ctx.Request, map[string]interface{}{
		"userId": userId,
	})
}

// SendMessage 发送消息到WebSocket连接
func SendMessage(s *melody.Session, message interface{}) error {
	data, err := json.Marshal(message)
	if err != nil {
		return err
	}
	return s.Write(data)
}

// SendErrorMessage 发送错误消息
func SendErrorMessage(s *melody.Session, errorMsg string) {
	SendMessage(s, NewErrorEvent(errorMsg))
}

// BroadcastToUser 向指定用户广播消息
func BroadcastToUser(userID uint, message interface{}) error {
	messageBytes, err := json.Marshal(message)
	if err != nil {
		return err
	}

	// 遍历所有连接，找到对应用户的连接
	return global.MELODY.BroadcastFilter(messageBytes, func(s *melody.Session) bool {
		if userIdInterface, exists := s.Get("userId"); exists {
			if userId, ok := userIdInterface.(uint); ok {
				return userId == userID
			}
		}
		return false
	})
}

// BroadcastMessageToUser 向指定用户广播流式消息
func BroadcastMessageToUser(userID uint, delta string) error {
	return BroadcastToUser(userID, NewMessageEvent(delta))
}

// BroadcastErrorToUser 向指定用户广播错误消息
func BroadcastErrorToUser(userID uint, errorMsg string) error {
	return BroadcastToUser(userID, NewErrorEvent(errorMsg))
}

// BroadcastDoneToUser 向指定用户广播完成消息
func BroadcastDoneToUser(userID uint, message, sessionId string) error {
	return BroadcastToUser(userID, NewDoneEvent(message, sessionId))
}

// GetConnectedUserCount 获取当前连接的用户数量
func GetConnectedUserCount() int {
	count := 0
	global.MELODY.BroadcastFilter([]byte{}, func(s *melody.Session) bool {
		if _, exists := s.Get("userId"); exists {
			count++
		}
		return false // 不实际发送消息，只统计
	})
	return count
}

// SendPingToUser 向指定用户发送心跳消息
func SendPingToUser(userID uint) error {
	pingMsg := Message{
		Type: "ping",
		Data: map[string]interface{}{
			"timestamp": time.Now().Unix(),
		},
	}
	return BroadcastToUser(userID, pingMsg)
}

// IsUserConnected 检查用户是否在线
func IsUserConnected(userID uint) bool {
	connected := false
	global.MELODY.BroadcastFilter([]byte{}, func(s *melody.Session) bool {
		if userIdInterface, exists := s.Get("userId"); exists {
			if userId, ok := userIdInterface.(uint); ok && userId == userID {
				connected = true
			}
		}
		return false // 不实际发送消息，只检查
	})
	return connected
}