3588AdminBackend/internal/service/registry.go

package service

import (
	"encoding/json"
	"strings"
	"sync"
	"time"

	"3588AdminBackend/internal/config"
	"3588AdminBackend/internal/models"
)

type RegistryService struct {
	cfg       *config.Config
	agent     *AgentClient
	discovery *DiscoveryService
	repo      DeviceRepository
	mu        sync.RWMutex
	devices   map[string]*models.Device
}

type DeviceRepository interface {
	Upsert(dev *models.Device) error
	List() ([]*models.Device, error)
}

func NewRegistryService(cfg *config.Config, agent *AgentClient, repo ...DeviceRepository) *RegistryService {
	var deviceRepo DeviceRepository
	if len(repo) > 0 {
		deviceRepo = repo[0]
	}
	s := &RegistryService{
		cfg:     cfg,
		agent:   agent,
		repo:    deviceRepo,
		devices: make(map[string]*models.Device),
	}
	// Load persisted devices from DB so offline devices still appear in the list
	if deviceRepo != nil {
		if saved, err := deviceRepo.List(); err == nil {
			for _, dev := range saved {
				if dev == nil || dev.DeviceID == "" {
					continue
				}
				dev.Online = false // mark offline until discovery confirms
				s.devices[dev.DeviceID] = dev
			}
		}
	}
	go s.startPruning()
	if agent != nil {
		go s.startGraphPolling()
	}
	go s.startDiscoveryPolling()
	return s
}

func (s *RegistryService) SetDiscovery(d *DiscoveryService) {
	s.discovery = d
}

func (s *RegistryService) startDiscoveryPolling() {
	// 启动后等 3 秒再开始第一次发现，避免和 startup 的 UDP 冲突
	time.Sleep(3 * time.Second)
	// First immediate discovery
	s.runDiscovery()
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()
	for range ticker.C {
		s.runDiscovery()
	}
}

func (s *RegistryService) runDiscovery() {
	s.mu.RLock()
	d := s.discovery
	s.mu.RUnlock()
	if d == nil {
		return
	}
	d.SearchDefault()
}

func (s *RegistryService) startGraphPolling() {
	ticker := time.NewTicker(30 * time.Second) // Pull every 30s
	for range ticker.C {
		s.mu.RLock()
		var onlineDevices []*models.Device
		for _, dev := range s.devices {
			if dev.Online {
				onlineDevices = append(onlineDevices, dev)
			}
		}
		s.mu.RUnlock()

		for _, dev := range onlineDevices {
			data, _, err := s.agent.Do("GET", dev.IP, dev.AgentPort, "/v1/graphs", nil)
			if err == nil {
				var graphs interface{}
				if err := json.Unmarshal(data, &graphs); err == nil {
					s.mu.Lock()
					dev.Graphs = graphs
					s.mu.Unlock()
					s.persistDevice(dev)
				}
			}
		}
	}
}

func (s *RegistryService) UpdateDevice(dev *models.Device) {
	s.mu.Lock()
	defer s.mu.Unlock()

	dev.LastSeenMs = time.Now().UnixMilli()
	dev.Online = true
	if current, ok := s.devices[dev.DeviceID]; ok && strings.TrimSpace(current.DeviceAlias) != "" {
		dev.DeviceAlias = strings.TrimSpace(current.DeviceAlias)
	} else if s.repo != nil {
		if saved, err := s.repo.List(); err == nil {
			for _, item := range saved {
				if item != nil && item.DeviceID == dev.DeviceID && strings.TrimSpace(item.DeviceAlias) != "" {
					dev.DeviceAlias = strings.TrimSpace(item.DeviceAlias)
					break
				}
			}
		}
	}
	s.devices[dev.DeviceID] = dev
	s.persistDevice(dev)
}

func (s *RegistryService) SetDeviceAlias(deviceID string, alias string) error {
	s.mu.Lock()
	defer s.mu.Unlock()

	alias = strings.TrimSpace(alias)
	if dev, ok := s.devices[deviceID]; ok {
		dev.DeviceAlias = alias
		s.persistDevice(dev)
		return nil
	}
	s.persistDevice(&models.Device{DeviceID: deviceID, DeviceAlias: alias})
	return nil
}

func (s *RegistryService) GetDevices() []*models.Device {
	s.mu.RLock()
	defer s.mu.RUnlock()

	list := make([]*models.Device, 0, len(s.devices))
	for _, dev := range s.devices {
		list = append(list, dev)
	}
	return list
}

// TouchDevice updates the LastSeenMs for a device to keep it online
// when accessed via HTTP API (not just UDP discovery)
func (s *RegistryService) TouchDevice(deviceID string) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if dev, ok := s.devices[deviceID]; ok {
		dev.LastSeenMs = time.Now().UnixMilli()
		dev.Online = true
		s.persistDevice(dev)
	}
}

func (s *RegistryService) startPruning() {
	ticker := time.NewTicker(2 * time.Second)
	for range ticker.C {
		s.mu.Lock()
		now := time.Now().UnixMilli()
		for _, dev := range s.devices {
			if now-dev.LastSeenMs > int64(s.cfg.OfflineAfterMs) {
				dev.Online = false
			}
		}
		s.mu.Unlock()
	}
}

func (s *RegistryService) persistDevice(dev *models.Device) {
	if s == nil || s.repo == nil || dev == nil {
		return
	}
	_ = s.repo.Upsert(dev)
}