LokiStack/backend/services/collector.go

package services

import (
	"context"
	"fmt"
	"time"

	"lokideploy/models"
	"lokideploy/repository"
	"lokideploy/utils"
)

type CollectorService struct{}

func NewCollectorService() *CollectorService {
	return &CollectorService{}
}

type CreateCollectorRequest struct {
	Name             string
	Type             string
	LokiDeploymentID uint
	HostID           *uint
	ClusterID        *uint
	TargetDeployment string
	Namespace        string
	LogPath          string
	Labels           string
}

// fillAssociations 填充关联的 Loki / Host / Cluster（用于配置生成与展示）
func (s *CollectorService) fillAssociations(c *models.Collector) {
	db := repository.GetDB()
	db.RLock()
	defer db.RUnlock()
	for i := range db.Deployments {
		if db.Deployments[i].ID == c.LokiDeploymentID {
			d := db.Deployments[i]
			c.LokiDeployment = &d
			break
		}
	}
	if c.HostID != nil {
		for i := range db.Hosts {
			if db.Hosts[i].ID == *c.HostID {
				h := db.Hosts[i]
				c.Host = &h
				break
			}
		}
	}
	if c.ClusterID != nil {
		for i := range db.Clusters {
			if db.Clusters[i].ID == *c.ClusterID {
				cl := db.Clusters[i]
				c.Cluster = &cl
				break
			}
		}
	}
}

// generateConfigSnapshot 按类型生成配置快照
func (s *CollectorService) generateConfigSnapshot(c *models.Collector) string {
	endpoint := ""
	if c.LokiDeployment != nil {
		endpoint = c.LokiDeployment.Endpoint
	}
	switch c.Type {
	case "host-promtail":
		return PromtailConfig(endpoint, c.LogPath)
	case "k8s-daemonset":
		return PromtailHelmValues(endpoint)
	case "k8s-fluentbit-sidecar":
		host, port := parseLokiURL(endpoint)
		app := c.TargetDeployment
		ns := c.Namespace
		if ns == "" {
			ns = "default"
		}
		return FluentbitConfig(host, port, c.LogPath, app, ns)
	}
	return ""
}

func (s *CollectorService) Create(req *CreateCollectorRequest) (*models.Collector, error) {
	switch req.Type {
	case "host-promtail", "k8s-daemonset", "k8s-fluentbit-sidecar":
	default:
		return nil, fmt.Errorf("不支持的采集器类型: %s", req.Type)
	}
	if req.LokiDeploymentID == 0 {
		return nil, fmt.Errorf("请选择关联的 Loki 部署")
	}
	if req.Type == "host-promtail" && (req.HostID == nil || *req.HostID == 0) {
		return nil, fmt.Errorf("主机采集器需要选择目标主机")
	}
	if (req.Type == "k8s-daemonset" || req.Type == "k8s-fluentbit-sidecar") && (req.ClusterID == nil || *req.ClusterID == 0) {
		return nil, fmt.Errorf("K8s 采集器需要选择目标集群")
	}
	if req.Type == "k8s-fluentbit-sidecar" && req.TargetDeployment == "" {
		return nil, fmt.Errorf("Sidecar 采集器需要填写目标 Deployment 名称")
	}

	refType := "cluster"
	if req.Type == "host-promtail" {
		refType = "host"
	}

	collector := models.Collector{
		ID:               repository.GetDB().NextID("collectors"),
		Name:             req.Name,
		Type:             req.Type,
		LokiDeploymentID: req.LokiDeploymentID,
		RefType:          refType,
		HostID:           req.HostID,
		ClusterID:        req.ClusterID,
		TargetDeployment: req.TargetDeployment,
		Namespace:        req.Namespace,
		LogPath:          req.LogPath,
		Labels:           req.Labels,
		Status:           "pending",
		CreatedAt:        time.Now(),
		UpdatedAt:        time.Now(),
	}

	s.fillAssociations(&collector)
	if collector.LokiDeployment == nil {
		return nil, fmt.Errorf("关联的 Loki 部署不存在")
	}
	collector.Config = s.generateConfigSnapshot(&collector)

	db := repository.GetDB()
	db.Lock()
	db.Collectors = append(db.Collectors, collector)
	db.Unlock()
	db.RequestSave()

	go s.Install(collector.ID)
	return &collector, nil
}

func (s *CollectorService) Get(id uint) (*models.Collector, error) {
	db := repository.GetDB()
	db.RLock()
	var found *models.Collector
	for i := range db.Collectors {
		if db.Collectors[i].ID == id {
			c := db.Collectors[i]
			found = &c
			break
		}
	}
	db.RUnlock()
	if found == nil {
		return nil, utils.ErrNotFound
	}
	s.fillAssociations(found)
	return found, nil
}

func (s *CollectorService) List(page, pageSize int) ([]models.Collector, int64, error) {
	db := repository.GetDB()
	db.RLock()
	total := int64(len(db.Collectors))
	start := (page - 1) * pageSize
	if start < 0 || start >= len(db.Collectors) {
		db.RUnlock()
		return []models.Collector{}, total, nil
	}
	end := start + pageSize
	if end > len(db.Collectors) {
		end = len(db.Collectors)
	}
	result := make([]models.Collector, end-start)
	copy(result, db.Collectors[start:end])
	db.RUnlock()
	for i := range result {
		s.fillAssociations(&result[i])
	}
	return result, total, nil
}

func (s *CollectorService) Update(id uint, updates map[string]interface{}) error {
	db := repository.GetDB()
	db.Lock()
	defer db.Unlock()
	for i := range db.Collectors {
		if db.Collectors[i].ID == id {
			if v, ok := updates["name"].(string); ok && v != "" {
				db.Collectors[i].Name = v
			}
			if v, ok := updates["status"].(string); ok && v != "" {
				db.Collectors[i].Status = v
			}
			if v, ok := updates["statusMessage"].(string); ok {
				db.Collectors[i].StatusMessage = v
			}
			if v, ok := updates["logPath"].(string); ok {
				db.Collectors[i].LogPath = v
			}
			if v, ok := updates["labels"].(string); ok {
				db.Collectors[i].Labels = v
			}
			db.Collectors[i].UpdatedAt = time.Now()
			db.RequestSave()
			return nil
		}
	}
	return utils.ErrNotFound
}

func (s *CollectorService) Delete(id uint) error {
	collector, err := s.Get(id)
	if err != nil {
		return utils.ErrNotFound
	}
	db := repository.GetDB()
	db.Lock()
	for i, c := range db.Collectors {
		if c.ID == id {
			db.Collectors = append(db.Collectors[:i], db.Collectors[i+1:]...)
			break
		}
	}
	db.Unlock()
	db.RequestSave()
	// 异步卸载（用拷贝的 collector，避免删除后取不到）
	go func() {
		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
		defer cancel()
		NewCollectorDeployer(collector, collector.LokiDeployment, nil).Uninstall(ctx)
	}()
	return nil
}

// Install 异步部署采集器
func (s *CollectorService) Install(id uint) {
	collector, err := s.Get(id)
	if err != nil {
		return
	}
	if collector.LokiDeployment == nil {
		s.Update(id, map[string]interface{}{"status": "error", "statusMessage": "关联的 Loki 不存在"})
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()
	if err := NewCollectorDeployer(collector, collector.LokiDeployment, nil).Install(ctx); err != nil {
		s.Update(id, map[string]interface{}{"status": "error", "statusMessage": err.Error()})
		return
	}
	s.Update(id, map[string]interface{}{"status": "active", "statusMessage": "采集器已部署"})
}

// Uninstall 异步卸载采集器
func (s *CollectorService) Uninstall(id uint) {
	collector, err := s.Get(id)
	if err != nil {
		return
	}
	if collector.LokiDeployment == nil {
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	NewCollectorDeployer(collector, collector.LokiDeployment, nil).Uninstall(ctx)
	s.Update(id, map[string]interface{}{"status": "pending", "statusMessage": "采集器已卸载"})
}