LokiStack/backend/services/collector_deployer.go

package services

import (
	"context"
	"encoding/json"
	"fmt"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"strings"
	"time"

	"lokistack/models"
)

// CollectorDeployer 采集器部署器，按 type 分发到三种安装/卸载流程
type CollectorDeployer struct {
	collector *models.Collector
	loki      *models.LokiDeployment
	logFunc   func(level, message, detail string)
}

func NewCollectorDeployer(c *models.Collector, l *models.LokiDeployment, logFunc func(level, message, detail string)) *CollectorDeployer {
	return &CollectorDeployer{collector: c, loki: l, logFunc: logFunc}
}

func (d *CollectorDeployer) log(level, message, detail string) {
	if d.logFunc != nil {
		d.logFunc(level, message, detail)
	}
}

// Install 按采集器类型分发安装
func (d *CollectorDeployer) Install(ctx context.Context) error {
	endpoint := d.loki.Endpoint
	if endpoint == "" {
		return fmt.Errorf("关联的 Loki 尚未部署，缺少推送地址（endpoint），请先成功部署该 Loki")
	}
	d.log("info", fmt.Sprintf("开始部署采集器 %s（类型 %s）", d.collector.Name, d.collector.Type), "")

	var err error
	switch d.collector.Type {
	case "host-promtail":
		err = d.installHost(ctx, endpoint)
	case "k8s-daemonset":
		err = d.installDaemonset(ctx, endpoint)
	case "k8s-fluentbit-sidecar":
		err = d.installSidecar(ctx, endpoint)
	default:
		return fmt.Errorf("不支持的采集器类型: %s", d.collector.Type)
	}
	if err != nil {
		d.log("error", fmt.Sprintf("采集器部署失败: %v", err), "")
		return err
	}
	d.log("success", "采集器部署完成", "")
	return nil
}

// Uninstall 按采集器类型分发卸载
func (d *CollectorDeployer) Uninstall(ctx context.Context) error {
	d.log("info", fmt.Sprintf("开始卸载采集器 %s（类型 %s）", d.collector.Name, d.collector.Type), "")
	switch d.collector.Type {
	case "host-promtail":
		return d.uninstallHost(ctx)
	case "k8s-daemonset":
		return d.uninstallDaemonset(ctx)
	case "k8s-fluentbit-sidecar":
		return d.uninstallSidecar(ctx)
	}
	return nil
}

// ===== host-promtail：SSH 安装 Promtail 二进制 + systemd =====

func (d *CollectorDeployer) installHost(ctx context.Context, endpoint string) error {
	host := d.collector.Host
	if host == nil {
		return fmt.Errorf("主机采集器缺少目标主机")
	}
	client, err := NewSSHClientWithAuth(HostSSHAuth(host))
	if err != nil {
		return err
	}
	defer client.Close()

	// 检测架构
	archOut, _, _ := client.Execute(ctx, "uname -m")
	arch := strings.TrimSpace(archOut)
	switch arch {
	case "x86_64":
		arch = "amd64"
	case "aarch64", "arm64":
		arch = "arm64"
	default:
		return fmt.Errorf("不支持的主机架构: %s", arch)
	}

	// 确定 Promtail 版本（优先使用关联的 Loki 版本）
	version := "3.0.0"
	if d.loki != nil && d.loki.Version != "" {
		version = d.loki.Version
	}

	// 创建必要目录
	d.log("info", "准备目录 /opt/promtail /etc/promtail /var/lib/promtail", "")
	if _, _, err := client.Execute(ctx, "mkdir -p /opt/promtail /etc/promtail /var/lib/promtail"); err != nil {
		return fmt.Errorf("创建 Promtail 目录失败: %w", err)
	}

	// 下载 Promtail
	d.log("info", fmt.Sprintf("下载 Promtail %s（%s）", version, arch), "")
	url := fmt.Sprintf("https://github.com/grafana/loki/releases/download/v%s/promtail-linux-%s.zip", version, arch)
	dl := fmt.Sprintf("cd /tmp && curl -fsSL -o p.zip '%s' && unzip -o p.zip -d /opt/promtail/ && chmod +x /opt/promtail/promtail-linux-%s", url, arch)
	if _, _, err := client.Execute(ctx, dl); err != nil {
		return fmt.Errorf("下载 Promtail 失败: %w", err)
	}
	client.Execute(ctx, "ln -sf /opt/promtail/promtail-linux-* /usr/local/bin/promtail 2>/dev/null || true")

	// 生成并上传配置
	logPath := d.collector.LogPath
	if logPath == "" {
		logPath = "/var/log/*.log"
	}
	conf := PromtailConfig(endpoint, logPath, "host", d.collector.Labels)
	d.log("info", "上传 Promtail 配置", "")
	if _, _, err := client.Execute(ctx, fmt.Sprintf("cat > /etc/promtail/promtail.yaml << 'EOF'\n%s\nEOF", conf)); err != nil {
		return fmt.Errorf("上传配置失败: %w", err)
	}

	// systemd 服务（生产级）
	service := `[Unit]
Description=Promtail Log Collector
Wants=network-online.target
After=network-online.target

[Service]
Type=simple
User=root
Group=root
ExecStart=/usr/local/bin/promtail -config.file=/etc/promtail/promtail.yaml -positions.enabled=true
ExecReload=/bin/kill -HUP $MAINPID
Restart=always
RestartSec=10
LimitNOFILE=65536
StandardOutput=append:/var/log/promtail.log
StandardError=append:/var/log/promtail.log

[Install]
WantedBy=multi-user.target
`
	d.log("info", "安装 systemd 服务", "")
	if _, _, err := client.Execute(ctx, fmt.Sprintf("cat > /etc/systemd/system/promtail.service << 'EOF'\n%s\nEOF", service)); err != nil {
		return fmt.Errorf("创建 systemd 服务失败: %w", err)
	}
	if _, _, err := client.Execute(ctx, "systemctl daemon-reload && systemctl enable promtail && systemctl restart promtail"); err != nil {
		return fmt.Errorf("启动 Promtail 失败: %w", err)
	}

	// 保存配置快照
	d.collector.Config = conf
	d.log("success", "Promtail 部署完成", "")
	return nil
}

func (d *CollectorDeployer) uninstallHost(ctx context.Context) error {
	host := d.collector.Host
	if host == nil {
		return nil
	}
	client, err := NewSSHClientWithAuth(HostSSHAuth(host))
	if err != nil {
		return err
	}
	defer client.Close()
	client.Execute(ctx, "systemctl stop promtail 2>/dev/null || true")
	client.Execute(ctx, "systemctl disable promtail 2>/dev/null || true")
	client.Execute(ctx, "rm -rf /etc/systemd/system/promtail.service /usr/local/bin/promtail /etc/promtail /opt/promtail /var/lib/promtail")
	client.Execute(ctx, "systemctl daemon-reload")
	return nil
}

// ===== k8s-daemonset：以原生清单安装 Promtail DaemonSet =====

func (d *CollectorDeployer) installDaemonset(ctx context.Context, endpoint string) error {
	cluster := d.collector.Cluster
	if cluster == nil {
		return fmt.Errorf("集群采集器缺少目标集群")
	}

	// 默认部署到关联 Loki 的命名空间(统一 log 命名空间策略)
	ns := d.collector.Namespace
	if ns == "" && d.loki != nil {
		ns = d.loki.Namespace
	}
	if ns == "" {
		ns = "log"
	}

	clusterName := "unknown"
	if cluster.Name != "" {
		clusterName = cluster.Name
	}

	manifest := PromtailDaemonsetYAML(ns, endpoint, clusterName, d.collector.Labels,
		d.collector.TargetNamespaces, d.collector.AppLabelSource, d.collector.AppInclude)

	// 关键顺序：先本地校验清单，再删旧资源。
	// 否则模板有误时会出现「旧 DaemonSet 已删、新的建不出来」→ 日志采集直接中断。
	if err := validateManifest(manifest); err != nil {
		d.log("error", fmt.Sprintf("生成的 Promtail 清单无效，已中止且未改动现有资源: %v", err), "")
		return fmt.Errorf("Promtail 清单校验失败（现有采集器未受影响）: %w", err)
	}

	// 删除历史残留并通过 Kubernetes API 提交资源。
	for _, kind := range []string{"daemonset", "configmap", "serviceaccount"} {
		_ = deleteK8sResource(ctx, cluster, kind, ns, "promtail")
	}
	if err := d.applyYAML(ctx, nil, manifest, "promtail"); err != nil {
		return err
	}
	d.log("info", fmt.Sprintf("等待 DaemonSet promtail 就绪(命名空间 %s)...", ns), "")
	if err := waitDaemonSet(ctx, cluster, ns, "promtail"); err != nil {
		return fmt.Errorf("Promtail DaemonSet 未就绪: %w", err)
	}
	svcURL := endpoint

	d.collector.Config = svcURL
	d.log("success", fmt.Sprintf("Promtail DaemonSet 已部署到 %s(推送至 %s)", ns, svcURL), "")
	return nil
}

// applyYAML 通过 API Server 应用采集器清单。
func (d *CollectorDeployer) applyYAML(ctx context.Context, _ []string, manifest, label string) error {
	if d.collector.Cluster == nil {
		return fmt.Errorf("采集器缺少集群")
	}
	if err := applyManifest(ctx, d.collector.Cluster, manifest); err != nil {
		d.log("error", fmt.Sprintf("应用 %s 失败: %v", label, err), "")
		return err
	}
	return nil
}

func (d *CollectorDeployer) uninstallDaemonset(ctx context.Context) error {
	cluster := d.collector.Cluster
	if cluster == nil {
		return nil
	}

	ns := d.collector.Namespace
	if ns == "" && d.loki != nil {
		ns = d.loki.Namespace
	}
	if ns == "" {
		ns = "log"
	}
	// clusterrole / clusterrolebinding 用 loki-promtail；
	// 数组末两项是历史遗留命名，一并清理避免升级后留下孤儿资源。
	for _, item := range [][2]string{
		{"daemonset", "promtail"}, {"configmap", "promtail"}, {"serviceaccount", "promtail"},
		{"clusterrole", "loki-promtail"}, {"clusterrolebinding", "loki-promtail"},
		{"clusterrole", "opsxloki-promtail"}, {"clusterrolebinding", "opsxloki-promtail"},
	} {
		if err := deleteK8sResource(ctx, cluster, item[0], ns, item[1]); err != nil {
			return err
		}
	}
	return nil
}

// ===== k8s-fluentbit-sidecar：kubectl patch 注入 Fluent Bit sidecar =====

func (d *CollectorDeployer) installSidecar(ctx context.Context, endpoint string) error {
	cluster := d.collector.Cluster
	if cluster == nil {
		return fmt.Errorf("Sidecar 采集器缺少目标集群")
	}
	target := d.collector.TargetDeployment
	if target == "" {
		return fmt.Errorf("Sidecar 采集器缺少目标 Deployment")
	}
	client, err := clusterClient(cluster, 30*time.Second)
	if err != nil {
		return err
	}

	ns := d.collector.Namespace
	if ns == "" {
		ns = "default"
	}
	logPath := d.collector.LogPath
	if logPath == "" {
		logPath = "/var/log/app/*.log"
	}
	app := target

	// 1. 检查目标 deployment 是否存在
	d.log("info", fmt.Sprintf("检查目标 Deployment %s/%s", ns, target), "")
	deployment, err := client.AppsV1().Deployments(ns).Get(ctx, target, metav1.GetOptions{})
	if err != nil {
		return fmt.Errorf("读取目标 Deployment %s/%s: %w", ns, target, err)
	}
	deploymentJSON, err := json.Marshal(deployment)
	if err != nil {
		return err
	}
	mount, err := resolveSidecarLogMount(deploymentJSON, target, logPath, d.collector.ID)
	if err != nil {
		return err
	}
	hasSidecar := false
	for _, container := range deployment.Spec.Template.Spec.Containers {
		if container.Name == "fluent-bit" {
			hasSidecar = true
			break
		}
	}

	// 3. 部署 ConfigMap
	lokiHost, lokiPort := parseLokiURL(endpoint)
	conf := FluentbitConfig(lokiHost, lokiPort, logPath, app, ns)
	configMapName := fmt.Sprintf("fluent-bit-config-%d", d.collector.ID)
	cm := FluentbitConfigMapYAML(conf, ns, configMapName)
	d.log("info", "部署 Fluent Bit ConfigMap", "")
	if err := applyManifest(ctx, cluster, cm); err != nil {
		return fmt.Errorf("ConfigMap 部署失败: %w", err)
	}
	d.log("info", fmt.Sprintf("为 %s 配置 Fluent Bit Sidecar", target), "")
	if err := patchDeployment(ctx, cluster, ns, target, SidecarPatchYAML(configMapName, mount)); err != nil {
		return fmt.Errorf("Sidecar 注入失败: %w", err)
	}
	if hasSidecar {
		if err := restartDeployment(ctx, cluster, ns, target); err != nil {
			return fmt.Errorf("Sidecar 配置已更新，但重启 Deployment 失败: %w", err)
		}
	}
	d.log("info", "等待 Pod 重建", "")
	if err := waitDeployment(ctx, cluster, ns, target); err != nil {
		return fmt.Errorf("Sidecar Deployment 未就绪: %w", err)
	}

	// DNS 自愈:CoreDNS 异常集群上 svc 域名不可解析 → 用 Loki ClusterIP 重建 ConfigMap 并滚动重启
	if strings.Contains(endpoint, ".svc") {
		host := strings.Split(strings.Split(endpoint, "//")[1], ":")[0]
		if podDNSBroken(ctx, cluster, ns, target, "fluent-bit", host) {
			d.log("warn", "fluent-bit 无法解析 Loki 服务名，回退 ClusterIP…", "")
			if ip, ierr := serviceClusterIP(ctx, cluster, lokiNamespaceOf(d.loki, endpoint), "loki"); ierr == nil && ip != "" {
				conf2 := FluentbitConfig(ip, lokiPortOf(endpoint), logPath, target, ns)
				cm2 := FluentbitConfigMapYAML(conf2, ns, configMapName)
				if err := applyManifest(ctx, cluster, cm2); err == nil {
					_ = restartDeployment(ctx, cluster, ns, target)
					d.log("info", "fluent-bit 推送地址已回退为 ClusterIP: "+ip, "")
				}
			}
		}
	}

	return nil
}

func (d *CollectorDeployer) uninstallSidecar(ctx context.Context) error {
	cluster := d.collector.Cluster
	if cluster == nil || d.collector.TargetDeployment == "" {
		return nil
	}
	client, err := clusterClient(cluster, 30*time.Second)
	if err != nil {
		return err
	}

	ns := d.collector.Namespace
	if ns == "" {
		ns = "default"
	}
	deployment, err := client.AppsV1().Deployments(ns).Get(ctx, d.collector.TargetDeployment, metav1.GetOptions{})
	if err != nil {
		return err
	}
	deploymentJSON, err := json.Marshal(deployment)
	if err != nil {
		return err
	}
	removePatch, err := sidecarRemovePatchYAML(deploymentJSON, d.collector.ID)
	if err != nil {
		return err
	}
	if err := patchDeployment(ctx, cluster, ns, d.collector.TargetDeployment, removePatch); err != nil {
		return fmt.Errorf("移除 Sidecar 失败: %w", err)
	}
	if err := deleteK8sResource(ctx, cluster, "configmap", ns, fmt.Sprintf("fluent-bit-config-%d", d.collector.ID)); err != nil {
		return err
	}

	return nil
}

// lokiPortOf 从 endpoint 提取端口(默认 3100)
func lokiPortOf(endpoint string) string {
	rest := strings.SplitN(endpoint, "//", 2)
	if len(rest) != 2 {
		return "3100"
	}
	hostPart := rest[1]
	if i := strings.Index(hostPart, ":"); i >= 0 {
		port := hostPart[i+1:]
		if j := strings.IndexAny(port, "/"); j >= 0 {
			port = port[:j]
		}
		if port != "" {
			return port
		}
	}
	return "3100"
}