SingBox-Gopanel/internal/service/device_state.go

package service

import (
	"fmt"
	"sort"
	"strings"
	"time"
	"xboard-go/internal/database"
	"xboard-go/internal/model"
)

const deviceStateTTL = 10 * time.Minute

type userDevicesSnapshot map[string][]string

func SaveUserNodeDevices(userID, nodeID int, ips []string) error {
	unique := make([]string, 0, len(ips))
	seen := make(map[string]struct{})
	for _, ip := range ips {
		trimmed := strings.TrimSpace(ip)
		if trimmed == "" {
			continue
		}
		if _, ok := seen[trimmed]; ok {
			continue
		}
		seen[trimmed] = struct{}{}
		unique = append(unique, trimmed)
	}
	sort.Strings(unique)

	if err := database.CacheSet(deviceStateKey(userID, nodeID), unique, deviceStateTTL); err != nil {
		return err
	}
	return syncUserOnlineDevices(userID, nodeID, unique)
}

func GetUsersDevices(userIDs []int) map[int][]string {
	result := make(map[int][]string, len(userIDs))

	for _, userID := range userIDs {
		snapshot, ok := database.CacheGetJSON[userDevicesSnapshot](deviceStateUserIndexKey(userID))
		if !ok {
			result[userID] = []string{}
			continue
		}

		merged := make([]string, 0)
		seen := make(map[string]struct{})
		for _, ips := range snapshot {
			for _, ip := range ips {
				if _, exists := seen[ip]; exists {
					continue
				}
				seen[ip] = struct{}{}
				merged = append(merged, ip)
			}
		}
		sort.Strings(merged)
		if len(merged) == 0 {
			merged = loadUserOnlineDevicesFromDB(userID)
		}
		result[userID] = merged
	}

	return result
}

func SetDevices(userID, nodeID int, ips []string) error {
	if err := SaveUserNodeDevices(userID, nodeID, ips); err != nil {
		return err
	}

	indexKey := deviceStateUserIndexKey(userID)
	indexSnapshot, _ := database.CacheGetJSON[userDevicesSnapshot](indexKey)
	if indexSnapshot == nil {
		indexSnapshot = make(userDevicesSnapshot)
	}
	indexSnapshot[fmt.Sprintf("%d", nodeID)] = normalizeIPs(ips)
	return database.CacheSet(indexKey, indexSnapshot, deviceStateTTL)
}

func normalizeIPs(ips []string) []string {
	seen := make(map[string]struct{})
	result := make([]string, 0, len(ips))
	for _, ip := range ips {
		ip = strings.TrimSpace(ip)
		if ip == "" {
			continue
		}
		if _, ok := seen[ip]; ok {
			continue
		}
		seen[ip] = struct{}{}
		result = append(result, ip)
	}
	sort.Strings(result)
	return result
}

func deviceStateKey(userID, nodeID int) string {
	return fmt.Sprintf("device_state:user:%d:node:%d", userID, nodeID)
}

func deviceStateUserIndexKey(userID int) string {
	return fmt.Sprintf("device_state:user:%d:index", userID)
}

func syncUserOnlineDevices(userID, nodeID int, ips []string) error {
	now := time.Now().Unix()
	expiresAt := now + int64(deviceStateTTL.Seconds())

	if err := database.DB.Where("user_id = ? AND node_id = ? AND expires_at <= ?", userID, nodeID, now).
		Delete(&model.UserOnlineDevice{}).Error; err != nil {
		return err
	}

	existing := make([]model.UserOnlineDevice, 0)
	if err := database.DB.Where("user_id = ? AND node_id = ?", userID, nodeID).Find(&existing).Error; err != nil {
		return err
	}

	existingByIP := make(map[string]model.UserOnlineDevice, len(existing))
	for _, item := range existing {
		existingByIP[item.IP] = item
	}

	seen := make(map[string]struct{}, len(ips))
	for _, ip := range ips {
		seen[ip] = struct{}{}
		if current, ok := existingByIP[ip]; ok {
			if err := database.DB.Model(&model.UserOnlineDevice{}).
				Where("id = ?", current.ID).
				Updates(map[string]any{
					"last_seen_at": now,
					"expires_at":   expiresAt,
					"updated_at":   now,
				}).Error; err != nil {
				return err
			}
			continue
		}

		record := model.UserOnlineDevice{
			UserID:      userID,
			NodeID:      nodeID,
			IP:          ip,
			FirstSeenAt: now,
			LastSeenAt:  now,
			ExpiresAt:   expiresAt,
			CreatedAt:   now,
			UpdatedAt:   now,
		}
		if err := database.DB.Create(&record).Error; err != nil {
			return err
		}
	}

	if len(existing) > 0 {
		staleIDs := make([]uint64, 0)
		for _, item := range existing {
			if _, ok := seen[item.IP]; !ok {
				staleIDs = append(staleIDs, item.ID)
			}
		}
		if len(staleIDs) > 0 {
			if err := database.DB.Where("id IN ?", staleIDs).Delete(&model.UserOnlineDevice{}).Error; err != nil {
				return err
			}
		}
	}

	var onlineCount int64
	if err := database.DB.Model(&model.UserOnlineDevice{}).
		Where("user_id = ? AND expires_at > ?", userID, now).
		Distinct("ip").
		Count(&onlineCount).Error; err == nil {
		count := int(onlineCount)
		_ = database.DB.Model(&model.User{}).Where("id = ?", userID).Updates(map[string]any{
			"online_count":   count,
			"last_online_at": time.Unix(now, 0),
			"updated_at":     now,
		}).Error
	}

	return nil
}

func loadUserOnlineDevicesFromDB(userID int) []string {
	now := time.Now().Unix()
	var records []model.UserOnlineDevice
	if err := database.DB.
		Select("ip").
		Where("user_id = ? AND expires_at > ?", userID, now).
		Order("ip ASC").
		Find(&records).Error; err != nil {
		return []string{}
	}

	seen := make(map[string]struct{}, len(records))
	ips := make([]string, 0, len(records))
	for _, item := range records {
		if item.IP == "" {
			continue
		}
		if _, ok := seen[item.IP]; ok {
			continue
		}
		seen[item.IP] = struct{}{}
		ips = append(ips, item.IP)
	}
	return ips
}