LokiStack/backend/services/collector.go

package services

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

	"lokistack/models"
	"lokistack/repository"
	"lokistack/utils"
)

type CollectorService struct {
	// opsMu 保护 opsInflight:同一采集器的安装/卸载互斥去重。
	// 采集器是 fire-and-forget(接口立即返回),没有去重时连点「部署」
	// 会对同一主机并发 docker rm -f + docker run,容器名冲突/配置互踩。
	opsMu       sync.Mutex
	opsInflight map[uint]bool
}

func NewCollectorService() *CollectorService {
	return &CollectorService{opsInflight: map[uint]bool{}}
}

// beginOp 声明开始一次安装/卸载;已在进行返回 false(直接忽略重复触发)
func (s *CollectorService) beginOp(id uint) bool {
	s.opsMu.Lock()
	defer s.opsMu.Unlock()
	if s.opsInflight == nil {
		s.opsInflight = map[uint]bool{}
	}
	if s.opsInflight[id] {
		return false
	}
	s.opsInflight[id] = true
	return true
}

func (s *CollectorService) endOp(id uint) {
	s.opsMu.Lock()
	delete(s.opsInflight, id)
	s.opsMu.Unlock()
}

type CreateCollectorRequest struct {
	Name             string
	Type             string
	LokiDeploymentID uint
	HostID           *uint
	ClusterID        *uint
	TargetDeployment string
	Namespace        string
	LogPath          string
	Labels           string
	ScrapeJobs       []models.ScrapeJob
	TargetNamespaces []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
			}
		}
	}
	// 配置快照按当前字段实时渲染：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) {
	// 采集器变更可能带来新的 app 标签（主机模式），下拉候选值必须失效
	defer InvalidateLabelCache()
	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 名称")
	}

	// 入库前统一校验采集项(含敏感路径拦截):Create 也要挡,
	// 否则坏数据先落库、部署时才报错(Update 已有校验)
	if err := models.ValidateScrapeJobs(req.ScrapeJobs); err != nil {
		return nil, err
	}
	if strings.TrimSpace(req.LogPath) != "" {
		if err := models.ValidateScrapeJobs([]models.ScrapeJob{{Name: "app", Paths: strings.Split(req.LogPath, ",")}}); err != nil {
			return nil, err
		}
	}

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

	// DaemonSet 的部署命名空间固定为 Loki 所在命名空间（通常 log）：
	// 采集器与 LokiStack 放在一起，用户不需要也不应该选。
	// 这里强制清空，渲染与部署两侧都会回退到 d.loki.Namespace。
	if req.Type == "k8s-daemonset" {
		req.Namespace = ""
	}

	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,
		TargetNamespaces: req.TargetNamespaces,
		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 {
	defer InvalidateLabelCache()
	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
			}
			// DaemonSet 的 namespace 一律不接受修改（固定跟随 Loki 所在命名空间）
			if v, ok := updates["namespace"].(string); ok && v != "" && db.Collectors[i].Type != "k8s-daemonset" {
				db.Collectors[i].Namespace = v
				db.Collectors[i].StatusMessage = "部署命名空间已保存，点击部署以应用新配置"
			}
			// 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
}

// UninstallCollectorsForDeployment 卸载并删除某部署下的所有采集器。
// 卸载失败时保留记录，避免远端资源失去跟踪。
func UninstallCollectorsForDeployment(ctx context.Context, deploymentID uint) {
	db := repository.GetDB()
	db.RLock()
	var targets []models.Collector
	for _, c := range db.Collectors {
		if c.LokiDeploymentID == deploymentID {
			targets = append(targets, c)
		}
	}
	db.RUnlock()

	for i := range targets {
		c := targets[i]
		if err := NewCollectorDeployer(&c, c.LokiDeployment, nil).Uninstall(ctx); err != nil {
			log.Printf("[deploy-delete] 采集器 %q 卸载失败，保留记录: %v", c.Name, err)
			continue
		}
		db.Lock()
		for j := range db.Collectors {
			if db.Collectors[j].ID == c.ID {
				db.Collectors = append(db.Collectors[:j], db.Collectors[j+1:]...)
				break
			}
		}
		db.RequestSave()
		db.Unlock()
	}
}

func (s *CollectorService) Delete(id uint) error {
	collector, err := s.Get(id)
	if err != nil {
		return utils.ErrNotFound
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
	defer cancel()
	// 先卸载远端资源，成功后再删记录；卸载失败直接返回，避免「界面删了、远端还在」
	if err := NewCollectorDeployer(collector, collector.LokiDeployment, nil).Uninstall(ctx); err != nil {
		return fmt.Errorf("卸载采集器失败，未删除记录: %w", err)
	}
	db := repository.GetDB()
	db.Lock()
	defer db.Unlock()
	for i, c := range db.Collectors {
		if c.ID == id {
			db.Collectors = append(db.Collectors[:i], db.Collectors[i+1:]...)
			db.RequestSave()
			return nil
		}
	}
	return utils.ErrNotFound
}

// Install 异步部署采集器
func (s *CollectorService) Install(id uint) {
	if !s.beginOp(id) {
		return // 同一采集器的安装/卸载正在进行,忽略重复触发
	}
	defer s.endOp(id)
	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) {
	if !s.beginOp(id) {
		return
	}
	defer s.endOp(id)
	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": ""})
}