LokiStack/backend/services/collector.go

package services

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"lokistack/models"
	"lokistack/repository"
	"lokistack/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
	ScrapeJobs       []models.ScrapeJob
}

// 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
			}
		}
	}
	// 配置快照按当前字段实时渲染：DB 里存的 Config 只是上次安装时的副本，
	// 用户改完采集项还没点「部署」时，看配置应该看到「将要生效的 YAML」。
	// renderConfig 只读 c 的字段、不取锁，所以在 RLock 内调用是安全的。
	if cfg, err := s.renderConfig(c); err == nil {
		c.Config = cfg
	}
}

// parseScrapeJobs 把接口层传来的 scrapeJobs（[]interface{} / JSON 字符串）解析成模型。
// gin 的 map 绑定拿不到强类型，统一过一遍 JSON 转换，避免手写一堆类型断言。
func parseScrapeJobs(v interface{}) ([]models.ScrapeJob, error) {
	if v == nil {
		return nil, nil
	}
	raw, err := json.Marshal(v)
	if err != nil {
		return nil, fmt.Errorf("采集项格式不合法：%w", err)
	}
	var jobs []models.ScrapeJob
	if err := json.Unmarshal(raw, &jobs); err != nil {
		return nil, fmt.Errorf("采集项格式不合法：%w", err)
	}
	return jobs, nil
}

// renderConfig 按类型渲染当前配置（不落库）。
func (s *CollectorService) renderConfig(c *models.Collector) (string, error) {
	endpoint := ""
	if c.LokiDeployment != nil {
		endpoint = c.LokiDeployment.Endpoint
	}
	clusterName := "host"
	if c.Cluster != nil && c.Cluster.Name != "" {
		clusterName = c.Cluster.Name
	}
	switch c.Type {
	case "host-promtail":
		// 主机名与 syslog 格式：优先用上次安装时探测到的真实值，
		// 从未安装过时退回平台登记的主机名/IP（预览会提示这一点）。
		hostname := c.ResolvedHostname
		if hostname == "" && c.Host != nil {
			hostname = c.Host.Name
		}
		if hostname == "" && c.Host != nil {
			hostname = c.Host.HostIP
		}
		// 预览要与线上一致：这里也走一遍推送地址纠正（0.0.0.0 / 容器内回环）
		pushEndpoint := endpoint
		hostIP := ""
		if c.Host != nil {
			hostIP = c.Host.HostIP
		}
		if v, perr := hostPushEndpoint(endpoint, hostIP, c.RunMode == HostRunModeDocker); perr == nil {
			pushEndpoint = v
		}
		return PromtailConfigForHost(HostPromtailOptions{
			Endpoint:       pushEndpoint,
			Cluster:        "host",
			LabelsJSON:     c.Labels,
			Hostname:       hostname,
			SyslogFormat:   c.SyslogFormat,
			SystemLogPaths: c.SystemLogPaths,
			LogPath:        c.LogPath,
			Jobs:           c.ScrapeJobs,
		})
	case "k8s-daemonset":
		ns := c.Namespace
		if ns == "" && c.LokiDeployment != nil {
			ns = c.LokiDeployment.Namespace
		}
		if ns == "" {
			ns = "log"
		}
		return PromtailDaemonsetYAML(ns, endpoint, clusterName, c.Labels, c.TargetNamespaces, c.AppLabelSource, c.AppInclude), nil
	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), nil
	}
	return "", nil
}

// generateConfigSnapshot 按类型生成配置快照（渲染失败返回空串，不阻断保存）
func (s *CollectorService) generateConfigSnapshot(c *models.Collector) string {
	cfg, err := s.renderConfig(c)
	if err != nil {
		return ""
	}
	return cfg
}

// PreviewConfig 按「当前已保存配置 + 本次提交的改动」渲染 YAML 预览。
// 只渲染不落库：用户在编辑器里点「预览」时用，避免为了看一眼 YAML 就改线上采集。
func (s *CollectorService) PreviewConfig(id uint, updates map[string]interface{}) (map[string]interface{}, error) {
	c, err := s.Get(id)
	if err != nil {
		return nil, err
	}
	if v, ok := updates["logPath"].(string); ok {
		c.LogPath = v
	}
	if v, ok := updates["labels"].(string); ok {
		c.Labels = v
	}
	if v, ok := updates["scrapeJobs"]; ok {
		jobs, perr := parseScrapeJobs(v)
		if perr != nil {
			return nil, perr
		}
		c.ScrapeJobs = jobs
	}
	if c.Type == "host-promtail" {
		if err := models.ValidateScrapeJobs(c.ScrapeJobs); err != nil {
			return nil, err
		}
	}
	cfg, err := s.renderConfig(c)
	if err != nil {
		return nil, err
	}
	warnings := make([]string, 0, 2)
	if c.Type == "host-promtail" {
		if c.SyslogFormat == "" {
			warnings = append(warnings, "尚未成功安装过：系统日志时间戳格式会在安装时按目标机实际格式探测（ISO8601 / 传统 syslog）")
		}
		if c.ResolvedHostname == "" {
			warnings = append(warnings, "尚未成功安装过：host 标签将使用平台登记的主机名，实际安装时会替换为目标机真实 hostname")
		}
		// 推送地址填写了通配/回环地址时提前提示：这类地址采集器推不上去，
		// 表现是「容器/systemd 都 running，但 Loki 里一条日志都没有」。
		rawEndpoint := ""
		if c.LokiDeployment != nil {
			rawEndpoint = c.LokiDeployment.Endpoint
		}
		hostIP := ""
		if c.Host != nil {
			hostIP = c.Host.HostIP
		}
		if fixed, perr := hostPushEndpoint(rawEndpoint, hostIP, c.RunMode == HostRunModeDocker); perr == nil && fixed != rawEndpoint {
			warnings = append(warnings, fmt.Sprintf("Loki 推送地址 %s 不可用，部署时会自动改为 %s（建议修正 Loki 部署的监听地址）", rawEndpoint, fixed))
		}
	}
	return map[string]interface{}{"config": cfg, "warnings": warnings, "hostname": c.ResolvedHostname, "syslogFormat": c.SyslogFormat}, nil
}

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("Kubernetes 采集器需要选择目标集群")
	}
	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,
		ScrapeJobs:       req.ScrapeJobs,
		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 {
				if v != db.Collectors[i].LogPath {
					db.Collectors[i].StatusMessage = "采集路径已保存，点击部署以应用新配置"
				}
				db.Collectors[i].LogPath = v
			}
			if v, ok := updates["labels"].(string); ok {
				db.Collectors[i].Labels = v
			}
			// targetNamespaces: 采集哪些命名空间（[]string；空=全集群）
			if v, ok := updates["targetNamespaces"]; ok {
				var nss []string
				switch t := v.(type) {
				case []interface{}:
					for _, item := range t {
						if s, ok := item.(string); ok {
							if s = strings.TrimSpace(s); s != "" {
								nss = append(nss, s)
							}
						}
					}
				case []string:
					for _, s := range t {
						if s = strings.TrimSpace(s); s != "" {
							nss = append(nss, s)
						}
					}
				case string:
					// 兼容 "astralx,log" 这种逗号分隔写法
					for _, s := range strings.Split(t, ",") {
						if s = strings.TrimSpace(s); s != "" {
							nss = append(nss, s)
						}
					}
				}
				if len(nss) != len(db.Collectors[i].TargetNamespaces) {
					db.Collectors[i].StatusMessage = "采集命名空间已保存，点击部署以应用新配置"
				} else {
					for j := range nss {
						if nss[j] != db.Collectors[i].TargetNamespaces[j] {
							db.Collectors[i].StatusMessage = "采集命名空间已保存，点击部署以应用新配置"
							break
						}
					}
				}
				db.Collectors[i].TargetNamespaces = nss
			}
			// appLabelSource：app 标签取值来源（空=默认 controller_name 回退链）
			if v, ok := updates["appLabelSource"].(string); ok {
				if strings.TrimSpace(v) != db.Collectors[i].AppLabelSource {
					db.Collectors[i].StatusMessage = "app 标签来源已保存，点击部署以应用新配置"
				}
				db.Collectors[i].AppLabelSource = strings.TrimSpace(v)
			}
			// appInclude：只采集这些 app（白名单，空=不限）
			if v, ok := updates["appInclude"]; ok {
				var apps []string
				switch t := v.(type) {
				case []interface{}:
					for _, item := range t {
						if s, ok := item.(string); ok {
							if s = strings.TrimSpace(s); s != "" {
								apps = append(apps, s)
							}
						}
					}
				case []string:
					for _, s := range t {
						if s = strings.TrimSpace(s); s != "" {
							apps = append(apps, s)
						}
					}
				case string:
					for _, s := range strings.Split(t, ",") {
						if s = strings.TrimSpace(s); s != "" {
							apps = append(apps, s)
						}
					}
				}
				if len(apps) != len(db.Collectors[i].AppInclude) {
					db.Collectors[i].StatusMessage = "采集 app 白名单已保存，点击部署以应用新配置"
				}
				db.Collectors[i].AppInclude = apps
			}
			// scrapeJobs：主机模式的结构化采集项（B 档）
			if v, ok := updates["scrapeJobs"]; ok {
				jobs, jerr := parseScrapeJobs(v)
				if jerr != nil {
					return jerr
				}
				if verr := models.ValidateScrapeJobs(jobs); verr != nil {
					return verr
				}
				if len(jobs) != len(db.Collectors[i].ScrapeJobs) {
					db.Collectors[i].StatusMessage = "采集项已保存，点击部署以应用新配置"
				}
				db.Collectors[i].ScrapeJobs = jobs
			}
			// 以下三个字段由安装流程写回（探测结果 / 渲染快照），不在前端表单里
			if v, ok := updates["resolvedHostname"].(string); ok && v != "" {
				db.Collectors[i].ResolvedHostname = v
			}
			if v, ok := updates["syslogFormat"].(string); ok && v != "" {
				db.Collectors[i].SyslogFormat = v
			}
			// 系统日志路径也是探测结果（Ubuntu 与 RHEL 系不同），
			// 不落库的话「配置预览」会和线上实际采集路径不一致。
			if v, ok := updates["systemLogPaths"].([]string); ok && len(v) > 0 {
				db.Collectors[i].SystemLogPaths = v
			}
			// runMode：安装时探测出的托管方式（docker | systemd），卸载后清空。
			if v, ok := updates["runMode"].(string); ok {
				db.Collectors[i].RunMode = v
			}
			if v, ok := updates["config"].(string); ok && v != "" {
				db.Collectors[i].Config = 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
	}
	// 安装时探测到的真实 hostname / syslog 格式 + 实际生效的配置都要落库，
	// 否则「配置预览」与线上真实配置不一致（用户没法核对）。
	updates := map[string]interface{}{"status": "active", "statusMessage": "采集器已部署"}
	if collector.Config != "" {
		updates["config"] = collector.Config
	}
	if collector.ResolvedHostname != "" {
		updates["resolvedHostname"] = collector.ResolvedHostname
	}
	if collector.SyslogFormat != "" {
		updates["syslogFormat"] = collector.SyslogFormat
	}
	if len(collector.SystemLogPaths) > 0 {
		updates["systemLogPaths"] = collector.SystemLogPaths
	}
	// 托管方式（docker | systemd）由安装流程探测后写回，前端据此展示运行方式
	updates["runMode"] = collector.RunMode
	s.Update(id, updates)
}

// 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": "采集器已卸载", "runMode": ""})
}