// main.go
package main

import (
	"bytes"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/shirou/gopsutil/v3/process"
)

// MonitorConfig 表示對(duì)單個(gè)進(jìn)程的監(jiān)控配置
type MonitorConfig struct {
	Names       []string `json:"names"`        // 按進(jìn)程名匹配
	PIDs        []int32  `json:"pids"`         // 指定 pid
	CPUThreshold float64 `json:"cpu_threshold"` // 百分比，如 80.0
	MemThreshold float64 `json:"mem_threshold"` // MB，如 500.0
	RestartCmd   string  `json:"restart_cmd"`   // 告警時(shí)執(zhí)行的重啟命令（可空）
}

// AlertPayload 告警時(shí)發(fā)送的 JSON 結(jié)構(gòu)
type AlertPayload struct {
	Time      time.Time `json:"time"`
	Host      string    `json:"host"`
	Process   string    `json:"process"`
	PID       int32     `json:"pid"`
	CPU       float64   `json:"cpu_percent"`
	MemoryMB  float64   `json:"memory_mb"`
	Triggered string    `json:"triggered"`
	Msg       string    `json:"msg"`
}

func main() {
	// CLI 參數(shù)
	cfgFile := flag.String("config", "", "配置 JSON 文件（可選），與命令行參數(shù)組合使用")
	names := flag.String("names", "", "要監(jiān)控的進(jìn)程名，逗號(hào)分隔（例如: nginx,mysqld）")
	pids := flag.String("pids", "", "要監(jiān)控的 pid，逗號(hào)分隔（例如: 123,456）")
	interval := flag.Duration("interval", 5*time.Second, "采樣間隔")
	cpuTh := flag.Float64("cpu", 80.0, "默認(rèn) CPU 百分比閾值（%）")
	memTh := flag.Float64("mem", 500.0, "默認(rèn) 內(nèi)存閾值（MB）")
	webhook := flag.String("webhook", "", "告警 webhook URL（POST 接收 JSON）")
	restart := flag.String("restart", "", "全局重啟命令（可選，覆蓋 config 中 restart_cmd）")
	flag.Parse()

	// 解析配置
	var monitors []MonitorConfig
	if *cfgFile != "" {
		f, err := os.ReadFile(*cfgFile)
		if err != nil {
			log.Fatalf("讀取配置文件失敗: %v", err)
		}
		if err := json.Unmarshal(f, &monitors); err != nil {
			log.Fatalf("解析配置文件失敗: %v", err)
		}
	}

	// 命令行 args 補(bǔ)充單一配置（如果用戶沒傳配置文件）
	if len(monitors) == 0 && (*names != "" || *pids != "") {
		m := MonitorConfig{
			CPUThreshold: *cpuTh,
			MemThreshold: *memTh,
		}
		if *names != "" {
			for _, n := range strings.Split(*names, ",") {
				n = strings.TrimSpace(n)
				if n != "" {
					m.Names = append(m.Names, n)
				}
			}
		}
		if *pids != "" {
			for _, ps := range strings.Split(*pids, ",") {
				if s := strings.TrimSpace(ps); s != "" {
					id, err := strconv.Atoi(s)
					if err == nil {
						m.PIDs = append(m.PIDs, int32(id))
					}
				}
			}
		}
		if *restart != "" {
			m.RestartCmd = *restart
		}
		monitors = append(monitors, m)
	}

	if len(monitors) == 0 {
		log.Fatalln("沒有任何監(jiān)控配置。請(qǐng)通過 -config 或 -names/-pids 提供配置。")
	}

	hostname, _ := os.Hostname()
	ctx, cancel := context.WithCancel(context.Background())
	wg := &sync.WaitGroup{}

	// 信號(hào)優(yōu)雅退出
	sigc := make(chan os.Signal, 1)
	signal.Notify(sigc, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigc
		log.Println("收到退出信號(hào)，正在優(yōu)雅停止...")
		cancel()
	}()

	// 啟動(dòng)每個(gè)監(jiān)控項(xiàng)的 goroutine
	for idx, mc := range monitors {
		wg.Add(1)
		go func(id int, cfg MonitorConfig) {
			defer wg.Done()
			monitorLoop(ctx, id, cfg, *interval, *webhook, hostname)
		}(idx, mc)
	}

	// 等待退出
	wg.Wait()
	log.Println("procmon 已退出")
}

// monitorLoop 對(duì)單個(gè) MonitorConfig 進(jìn)行輪詢監(jiān)控
func monitorLoop(ctx context.Context, id int, cfg MonitorConfig, interval time.Duration, webhook, host string) {
	logPrefix := fmt.Sprintf("[monitor-%d] ", id)
	logger := log.New(os.Stdout, logPrefix, log.LstdFlags)

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	// 用于去重告警（避免短時(shí)間內(nèi)頻繁告警）
	alerted := make(map[int32]time.Time)
	alertCooldown := 30 * time.Second // 同一 pid 告警最小間隔

	for {
		select {
		case <-ctx.Done():
			logger.Println("停止監(jiān)控（context canceled）")
			return
		case <-ticker.C:
			procs, err := process.Processes()
			if err != nil {
				logger.Printf("獲取進(jìn)程列表失敗: %v\n", err)
				continue
			}
			now := time.Now()
			for _, p := range procs {
				match := false
				// 匹配 PID 列表
				for _, pid := range cfg.PIDs {
					if p.Pid == pid {
						match = true
						break
					}
				}
				// 匹配名字列表（如果未通過 pid 匹配）
				if !match && len(cfg.Names) > 0 {
					name, err := p.Name()
					if err == nil {
						for _, nm := range cfg.Names {
							if strings.EqualFold(name, nm) {
								match = true
								break
							}
						}
					}
				}
				if !match {
					continue
				}

				// 獲取 CPU & Mem
				// Percent 需要傳入一個(gè)間隔來計(jì)算；這里使用 0 來獲取自上次調(diào)用以后的值（某些平臺(tái)）
				// 更可靠的做法是調(diào)用 Percent(interval)；為了簡單與跨平臺(tái)，這里使用 Percent(0)
				cpuPercent, errCpu := p.CPUPercent()
				memInfo, errMem := p.MemoryInfo()
				if errCpu != nil || errMem != nil || memInfo == nil {
					// 有時(shí)權(quán)限原因無法讀取某些信息
					logger.Printf("讀取進(jìn)程 %d 信息失敗: cpuErr=%v memErr=%v\n", p.Pid, errCpu, errMem)
					continue
				}
				memMB := float64(memInfo.RSS) / 1024.0 / 1024.0

				// 打印日志
				name, _ := p.Name()
				logger.Printf("進(jìn)程 %s pid=%d cpu=%.2f%% mem=%.2fMB\n", name, p.Pid, cpuPercent, memMB)

				// 判斷閾值
				triggered := ""
				if cfg.CPUThreshold > 0 && cpuPercent >= cfg.CPUThreshold {
					triggered = "cpu"
				}
				if cfg.MemThreshold > 0 && memMB >= cfg.MemThreshold {
					if triggered == "" {
						triggered = "mem"
					} else {
						triggered = "cpu+mem"
					}
				}
				if triggered != "" {
					lastAlert, ok := alerted[p.Pid]
					if ok && now.Sub(lastAlert) < alertCooldown {
						// 跳過頻繁告警
						logger.Printf("已在 cooldown 中，跳過 pid=%d 的告警\n", p.Pid)
						continue
					}
					alerted[p.Pid] = now

					payload := AlertPayload{
						Time:      now,
						Host:      host,
						Process:   name,
						PID:       p.Pid,
						CPU:       cpuPercent,
						MemoryMB:  memMB,
						Triggered: triggered,
						Msg:       fmt.Sprintf("process %s (pid=%d) exceeded threshold (%s)", name, p.Pid, triggered),
					}
					// 本地日志告警
					logger.Printf("ALERT: %s\n", payload.Msg)

					// 發(fā)送 webhook（如果配置）
					if webhook != "" {
						go func(pl AlertPayload) {
							if err := postAlert(webhook, pl); err != nil {
								logger.Printf("發(fā)送 webhook 失敗: %v\n", err)
							} else {
								logger.Printf("告警已發(fā)送到 %s\n", webhook)
							}
						}(payload)
					}

					// 執(zhí)行重啟命令（config 中或全局傳入） —— 先嘗試 graceful terminate 再執(zhí)行重啟命令（如果提供）
					if cfg.RestartCmd != "" {
						go func(cmdStr string, targetPid int32) {
							logger.Printf("嘗試殺掉 pid=%d 并執(zhí)行重啟命令: %s\n", targetPid, cmdStr)
							// 發(fā)送 TERM
							_ = p.SendSignal(syscall.SIGTERM)
							// 等待短時(shí)間讓進(jìn)程退出
							time.Sleep(2 * time.Second)
							// 強(qiáng)制 kill 如果還存在
							exists, _ := process.PidExists(targetPid)
							if exists {
								_ = p.Kill()
							}
							// 執(zhí)行重啟命令（通過 shell）
							cmd := exec.Command("/bin/sh", "-c", cmdStr)
							out, err := cmd.CombinedOutput()
							if err != nil {
								logger.Printf("執(zhí)行重啟命令失敗: %v. output: %s\n", err, string(out))
							} else {
								logger.Printf("重啟命令已執(zhí)行, output: %s\n", string(out))
							}
						}(cfg.RestartCmd, p.Pid)
					}
				}
			} // end for procs
		} // end ticker select
	} // end for
}

// postAlert 以 JSON POST 方式發(fā)送告警
func postAlert(webhook string, payload AlertPayload) error {
	bs, _ := json.Marshal(payload)
	req, err := http.NewRequest("POST", webhook, bytes.NewReader(bs))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	client := &http.Client{Timeout: 8 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return fmt.Errorf("webhook 返回非 2xx: %s", resp.Status)
	}
	return nil
}

go mod init procmon
go get github.com/shirou/gopsutil/v3/process

[
  {
    "names": ["nginx"],
    "cpu_threshold": 70.0,
    "mem_threshold": 300,
    "restart_cmd": "systemctl restart nginx"
  },
  {
    "names": ["mysqld"],
    "cpu_threshold": 85.0,
    "mem_threshold": 2048,
    "restart_cmd": "systemctl restart mysql"
  }
]