aiops/backend/controller/agent_task.go

package controller

import (
	"fmt"
	"time"

	"ezops/middleware"
	"ezops/model"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"gorm.io/gorm"
)

// AgentTaskHandler Agent 远程任务管理 // Logan.Li
type AgentTaskHandler struct {
	db *gorm.DB
}

func NewAgentTaskHandler(db *gorm.DB) *AgentTaskHandler {
	return &AgentTaskHandler{db: db}
}

// CreateTask 下发远程任务（前端调用，需 JWT）
// POST /api/v1/agent/tasks
func (h *AgentTaskHandler) CreateTask(c *gin.Context) {
	var req struct {
		Hostname string `json:"hostname"`
		Type     string `json:"type"` // command / script
		Command  string `json:"command"`
		Timeout  int    `json:"timeout"`
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(400, model.Response{Code: 400, Message: "参数错误: " + err.Error()})
		return
	}
	if req.Hostname == "" || req.Command == "" {
		c.JSON(400, model.Response{Code: 400, Message: "hostname 和 command 不能为空"})
		return
	}
	if req.Timeout == 0 {
		req.Timeout = 60
	}
	if req.Type == "" {
		req.Type = "command"
	}

	username, _ := c.Get("username")

	task := model.AgentTask{
		TaskID:    uuid.New().String(),
		Hostname:  req.Hostname,
		Type:      req.Type,
		Command:   req.Command,
		Timeout:   req.Timeout,
		Status:    "pending",
		CreatedBy: fmt.Sprintf("%v", username),
	}

	if err := h.db.Create(&task).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: "创建失败: " + err.Error()})
		return
	}

	c.JSON(200, model.Response{Code: 200, Message: "ok", Data: task})
}

// PollTasks Agent 轮询拉取任务（Agent 调用，无需 JWT）
// GET /api/v1/agent/tasks/poll?hostname=xxx
func (h *AgentTaskHandler) PollTasks(c *gin.Context) {
	hostname := c.Query("hostname")
	if hostname == "" {
		c.JSON(400, model.Response{Code: 400, Message: "hostname 不能为空"})
		return
	}
	if !middleware.RequireAgentHostname(c, hostname) {
		return
	}

	// 查找该主机 pending 的任务，按时间排序，最多取 5 个
	var tasks []model.AgentTask
	h.db.Where("hostname = ? AND status = ?", hostname, "pending").
		Order("created_at ASC").
		Limit(5).
		Find(&tasks)

	// 标记为 running（防其他 agent 重复拉取）
	now := time.Now()
	for _, t := range tasks {
		h.db.Model(&model.AgentTask{}).Where("id = ? AND status = ?", t.ID, "pending").
			Updates(map[string]interface{}{"status": "running", "picked_at": now})
	}

	c.JSON(200, model.Response{Code: 200, Data: tasks})
}

// SubmitResult Agent 回传执行结果（Agent 调用，无需 JWT）
// POST /api/v1/agent/tasks/:task_id/result
func (h *AgentTaskHandler) SubmitResult(c *gin.Context) {
	taskID := c.Param("task_id")
	var task model.AgentTask
	if err := h.db.Where("task_id = ?", taskID).Take(&task).Error; err != nil {
		c.JSON(404, model.Response{Code: 404, Message: "任务不存在"})
		return
	}
	if !middleware.RequireAgentHostname(c, task.Hostname) {
		return
	}

	var req struct {
		ExitCode int    `json:"exit_code"`
		Result   string `json:"result"`
		Status   string `json:"status"` // success / failed / timeout
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(400, model.Response{Code: 400, Message: "参数错误"})
		return
	}

	now := time.Now()
	updates := map[string]interface{}{
		"exit_code":   req.ExitCode,
		"result":      req.Result,
		"status":      req.Status,
		"finished_at": now,
	}

	if err := h.db.Model(&model.AgentTask{}).Where("task_id = ?", taskID).Updates(updates).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: "更新失败"})
		return
	}

	c.JSON(200, model.Response{Code: 200, Message: "ok"})
}

// ListTasks 查询任务列表（前端调用，需 JWT）
// GET /api/v1/agent/tasks?hostname=xxx&status=xxx
func (h *AgentTaskHandler) ListTasks(c *gin.Context) {
	query := h.db.Model(&model.AgentTask{})

	if hostname := c.Query("hostname"); hostname != "" {
		query = query.Where("hostname = ?", hostname)
	}
	if status := c.Query("status"); status != "" {
		query = query.Where("status = ?", status)
	}

	var tasks []model.AgentTask
	query.Order("created_at DESC").Limit(200).Find(&tasks)

	c.JSON(200, model.Response{Code: 200, Data: tasks})
}

// GetTask 获取单个任务详情
// GET /api/v1/agent/tasks/:id
func (h *AgentTaskHandler) GetTask(c *gin.Context) {
	var task model.AgentTask
	if err := h.db.First(&task, c.Param("id")).Error; err != nil {
		c.JSON(404, model.Response{Code: 404, Message: "任务不存在"})
		return
	}
	c.JSON(200, model.Response{Code: 200, Data: task})
}

// ===== Agent 告警规则 =====

// QuickSetupRules POST /api/v1/agent/rules/quick-setup
// 一键创建默认告警规则（CPU/内存/磁盘/网络带宽），作用于所有在线主机
func (h *AgentTaskHandler) QuickSetupRules(c *gin.Context) {
	var req struct {
		Type string `json:"type"` // "ecs" 或留空
	}
	c.ShouldBindJSON(&req)

	// 获取所有在线主机
	var onlineHosts []string
	h.db.Model(&model.Server{}).
		Where("agent_status = ?", "online").
		Pluck("hostname", &onlineHosts)

	if len(onlineHosts) == 0 {
		// 如果没有在线主机，使用 UTC 边界从 Agent 报告中查找最近 3 分钟活跃主机。 // Logan.Li
		h.db.Model(&model.AgentReport{}).
			Where("report_time >= ?", time.Now().UTC().Add(-3*time.Minute)).
			Distinct("hostname").
			Pluck("hostname", &onlineHosts)
	}

	if len(onlineHosts) == 0 {
		c.JSON(400, model.Response{Code: 400, Message: "没有在线的主机，请先安装并启动 Agent"})
		return
	}

	type ruleDef struct {
		Name      string
		Metric    string
		Operator  string
		Threshold float64
		Duration  int
	}

	defaultRules := []ruleDef{
		{Name: "CPU使用率(Warn)", Metric: "cpu", Operator: ">", Threshold: 80},
		{Name: "内存使用率(Warn)", Metric: "mem", Operator: ">", Threshold: 80},
		{Name: "磁盘使用率(Warn)", Metric: "disk", Operator: ">", Threshold: 80},
		{Name: "FD使用率(Warn)", Metric: "fd", Operator: ">", Threshold: 80},
	}

	created := 0
	skipped := 0
	for _, host := range onlineHosts {
		for _, def := range defaultRules {
			// 检查是否已存在同名同主机的规则，避免重复创建
			var existing model.AgentRule
			h.db.Where("name = ? AND hostname = ?", def.Name, host).First(&existing)
			if existing.ID > 0 {
				skipped++
				continue
			}

			rule := model.AgentRule{
				Name:      def.Name,
				Hostname:  host,
				Metric:    def.Metric,
				Operator:  def.Operator,
				Threshold: def.Threshold,
				Enabled:   true,
			}
			if err := h.db.Create(&rule).Error; err != nil {
				continue
			}
			created++
		}
	}

	c.JSON(200, model.Response{
		Code:    200,
		Message: fmt.Sprintf("已为 %d 台主机创建 %d 条规则（跳过 %d 条已存在的)", len(onlineHosts), created, skipped),
		Data:    gin.H{"hosts": len(onlineHosts), "created": created, "skipped": skipped},
	})
}

// ListRules GET /api/v1/agent/rules
func (h *AgentTaskHandler) ListRules(c *gin.Context) {
	var rules []model.AgentRule
	h.db.Order("created_at DESC").Find(&rules)
	c.JSON(200, model.Response{Code: 200, Data: rules})
}

// CreateRule POST /api/v1/agent/rules
func (h *AgentTaskHandler) CreateRule(c *gin.Context) {
	var rule model.AgentRule
	if err := c.ShouldBindJSON(&rule); err != nil {
		c.JSON(400, model.Response{Code: 400, Message: err.Error()})
		return
	}
	rule.ID = 0
	rule.Enabled = true
	if err := h.db.Create(&rule).Error; err != nil {
		c.JSON(500, model.Response{Code: 500, Message: err.Error()})
		return
	}
	c.JSON(200, model.Response{Code: 200, Data: rule})
}

// UpdateRule PUT /api/v1/agent/rules/:id
func (h *AgentTaskHandler) UpdateRule(c *gin.Context) {
	var rule model.AgentRule
	if err := h.db.First(&rule, c.Param("id")).Error; err != nil {
		c.JSON(404, model.Response{Code: 404, Message: "规则不存在"})
		return
	}
	var update model.AgentRule
	if err := c.ShouldBindJSON(&update); err != nil {
		c.JSON(400, model.Response{Code: 400, Message: err.Error()})
		return
	}
	h.db.Model(&rule).Updates(update)
	c.JSON(200, model.Response{Code: 200, Data: rule})
}

// DeleteRule DELETE /api/v1/agent/rules/:id
func (h *AgentTaskHandler) DeleteRule(c *gin.Context) {
	h.db.Delete(&model.AgentRule{}, c.Param("id"))
	c.JSON(200, model.Response{Code: 200, Message: "ok"})
}

// ListAlerts GET /api/v1/agent/alerts
func (h *AgentTaskHandler) ListAlerts(c *gin.Context) {
	query := h.db.Model(&model.AgentAlert{})
	if c.Query("resolved") == "false" {
		query = query.Where("resolved = ?", false)
	}
	if hostname := c.Query("hostname"); hostname != "" {
		query = query.Where("hostname = ?", hostname)
	}
	var alerts []model.AgentAlert
	query.Order("created_at DESC").Limit(200).Find(&alerts)
	c.JSON(200, model.Response{Code: 200, Data: alerts})
}

// ResolveAlert POST /api/v1/agent/alerts/:id/resolve
func (h *AgentTaskHandler) ResolveAlert(c *gin.Context) {
	now := time.Now()
	h.db.Model(&model.AgentAlert{}).Where("id = ?", c.Param("id")).
		Updates(map[string]interface{}{"resolved": true, "resolved_at": now})
	c.JSON(200, model.Response{Code: 200, Message: "ok"})
}