star/nmon/mon.go

package nmon

import (
	"fmt"
	"os"
	"sync"
	"time"

	"github.com/gdamore/tcell/v2"
	"github.com/rivo/tview"
	"github.com/shirou/gopsutil/v4/cpu"
	"github.com/shirou/gopsutil/v4/disk"
	"github.com/shirou/gopsutil/v4/load"
	"github.com/shirou/gopsutil/v4/mem"
	"github.com/shirou/gopsutil/v4/net"
	"github.com/spf13/cobra"
)

var (
	refreshInterval time.Duration
	savePath        string
	app             *tview.Application
	showCores       bool
	noTUI           bool
	diskIOHistory   = make(map[string]disk.IOCountersStat)
	netStats        = make(map[string]*NetStat)
	fileMutex       sync.Mutex
)

type NetStat struct {
	LastBytesSent   uint64
	LastBytesRecv   uint64
	LastTime        time.Time
	CurrentSentRate float64
	CurrentRecvRate float64
}

var Cmd = &cobra.Command{
	Use:   "nmon",
	Short: "System Monitoring Tool",
	Run: func(cmd *cobra.Command, args []string) {
		if noTUI {
			runTextMode()
			return
		}
		startTUI()
	},
}

func init() {
	Cmd.Flags().BoolVarP(&noTUI, "no-tui", "t", false, "Run in text mode")
	Cmd.Flags().BoolVarP(&showCores, "cores", "c", false, "Show per-core CPU usage")
	Cmd.Flags().DurationVarP(&refreshInterval, "interval", "i", time.Second, "Refresh interval")
	Cmd.Flags().StringVarP(&savePath, "save", "s", "", "Save monitoring data to file (nmon format)")
}

func main() {
	if err := Cmd.Execute(); err != nil {
		fmt.Println(err)
	}
}

func createClickableTextView(title string) *tview.TextView {
	view := tview.NewTextView().SetDynamicColors(true)
	view.SetBorder(true).SetTitle(title)
	return view
}

func startTUI() {
	app = tview.NewApplication()
	flex := tview.NewFlex().SetDirection(tview.FlexRow)

	cpuView := tview.NewTextView().SetDynamicColors(true)
	memView := tview.NewTextView().SetDynamicColors(true)
	netView := tview.NewTextView().SetDynamicColors(true)
	diskView := tview.NewTextView().SetDynamicColors(true)

	// 动态布局更新函数
	ioStats, _ := disk.IOCounters()
	updateLayout := func() {
		flex.Clear()
		cpuLines := 2 // 固定基础2行（CPU Load + Load Avg）
		if showCores {
			perCPU, _ := cpu.Percent(0, true)
			cpuLines = 2 + len(perCPU)
		}
		flex.AddItem(cpuView, cpuLines, 1, false)
		flex.AddItem(memView, 3, 1, false)
		flex.AddItem(diskView, len(ioStats), 1, false)
		flex.AddItem(netView, 3, 1, false)
	}

	// 初始化数据
	updateLayout()

	app.SetMouseCapture(func(event *tcell.EventMouse, action tview.MouseAction) (*tcell.EventMouse, tview.MouseAction) {
		// 示例：打印点击坐标
		if action == tview.MouseLeftClick {
			showCores = !showCores
			app.QueueUpdateDraw(updateLayout)
		}
		return event, action
	})

	go updateLoop(cpuView, memView, netView, diskView)

	app.SetInputCapture(func(event *tcell.EventKey) *tcell.EventKey {
		if event.Key() == tcell.KeyCtrlC {
			app.Stop()
		}
		return event
	})

	if err := app.SetRoot(flex, true).Run(); err != nil {
		panic(err)
	}
}

func updateLoop(cpuView, memView, netView, diskView *tview.TextView) {
	var lastSave time.Time
	for {
		app.QueueUpdateDraw(func() {
			updateCPU(cpuView)
			updateMemory(memView)
			updateNetwork(netView)
			updateDiskIO(diskView)
		})

		if savePath != "" && time.Since(lastSave) > time.Minute {
			go saveNMONData()
			lastSave = time.Now()
		}

		time.Sleep(refreshInterval)
	}
}

func updateCPU(view *tview.TextView) {
	percent, _ := cpu.Percent(time.Second, false)
	loadAvg, _ := load.Avg()

	text := fmt.Sprintf("[red]CPU Load: [white]%.2f%%\n[red]Load Avg: [white]%.2f, %.2f, %.2f",
		percent[0],
		loadAvg.Load1,
		loadAvg.Load5,
		loadAvg.Load15)

	if showCores {
		perCPU, _ := cpu.Percent(time.Second, true)
		for i, p := range perCPU {
			text += fmt.Sprintf("\nCore %d: [green]%s[white] %.2f%%",
				i+1,
				progressBar(p, 20),
				p)
		}
	}
	view.SetText(text)
}

func updateMemory(view *tview.TextView) {
	memStat, _ := mem.VirtualMemory()
	text := fmt.Sprintf("[red]Memory: [white]%.2f%% Used (%.2f GB / %.2f GB)",
		memStat.UsedPercent,
		float64(memStat.Used)/1024/1024/1024,
		float64(memStat.Total)/1024/1024/1024)
	view.SetText(text + "\n" + progressBar(memStat.UsedPercent, 50))
}

func updateNetwork(view *tview.TextView) {
	interfaces, _ := net.IOCounters(true)
	now := time.Now()
	var totalSent, totalRecv float64

	for _, iface := range interfaces {
		if iface.Name == "lo" {
			continue
		}

		stat, exists := netStats[iface.Name]
		if !exists {
			stat = &NetStat{
				LastBytesSent: iface.BytesSent,
				LastBytesRecv: iface.BytesRecv,
				LastTime:      now,
			}
			netStats[iface.Name] = stat
			continue
		}

		timeDiff := now.Sub(stat.LastTime).Seconds()
		sentDiff := float64(iface.BytesSent - stat.LastBytesSent)
		recvDiff := float64(iface.BytesRecv - stat.LastBytesRecv)

		stat.CurrentSentRate = sentDiff / timeDiff
		stat.CurrentRecvRate = recvDiff / timeDiff
		stat.LastBytesSent = iface.BytesSent
		stat.LastBytesRecv = iface.BytesRecv
		stat.LastTime = now

		totalSent += stat.CurrentSentRate
		totalRecv += stat.CurrentRecvRate
	}

	view.SetText(fmt.Sprintf("[red]Network: [white]↑%s  ↓%s",
		formatSpeed(totalSent),
		formatSpeed(totalRecv)))
}

func updateDiskIO(view *tview.TextView) {
	ioStats, _ := disk.IOCounters()
	text := "[red]Disk I/O:\n"

	for name, io := range ioStats {
		last, exists := diskIOHistory[name]
		if exists {
			interval := float64(refreshInterval.Seconds())
			readSpeed := float64(io.ReadBytes-last.ReadBytes) / interval
			writeSpeed := float64(io.WriteBytes-last.WriteBytes) / interval

			text += fmt.Sprintf("  %-8s R:%s  W:%s\n",
				name,
				formatSpeed(readSpeed),
				formatSpeed(writeSpeed))
		}
		diskIOHistory[name] = io
	}
	view.SetText(text)
}

func progressBar(percent float64, width int) string {
	filled := int(percent / 100 * float64(width))
	bar := "[green]"
	for i := 0; i < width; i++ {
		if i < filled {
			bar += "■"
		} else {
			bar += "□"
		}
	}
	return bar + "[white]"
}

func formatSpeed(speed float64) string {
	const (
		KB = 1024
		MB = KB * 1024
	)

	switch {
	case speed >= MB:
		return fmt.Sprintf("%7.3f MB/s", speed/MB)
	case speed >= KB:
		return fmt.Sprintf("%7.3f KB/s", speed/KB)
	default:
		return fmt.Sprintf("%7.3f  B/s", speed)
	}
}

func saveNMONData() {
	fileMutex.Lock()
	defer fileMutex.Unlock()

	f, err := os.OpenFile(savePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return
	}
	defer f.Close()

	now := time.Now().Format("2006-01-02 15:04:05")
	cpuPercent, _ := cpu.Percent(0, false)
	memStat, _ := mem.VirtualMemory()
	ioStats, _ := disk.IOCounters()

	fmt.Fprintf(f, "%s,CPU_ALL,%.2f\n", now, cpuPercent[0])
	fmt.Fprintf(f, "%s,MEM,%.2f,%.2f\n", now, memStat.UsedPercent, memStat.Available)

	for name, io := range ioStats {
		fmt.Fprintf(f, "%s,DISKIO,%s,%d,%d\n",
			now, name, io.ReadBytes, io.WriteBytes)
	}
}

// 文本模式相关函数
func runTextMode() {

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

	for {
		fmt.Println(getCPUInfo())
		fmt.Println(getMemoryInfo())
		fmt.Println(getNetworkInfo())
		fmt.Println(getDiskIOInfo())

		if savePath != "" {
			go saveNMONData()
		}
		<-ticker.C

	}
}

// 数据获取函数
func getCPUInfo() string {

	percent, _ := cpu.Percent(time.Second, false)
	loadAvg, _ := load.Avg()
	fmt.Print("\033[H\033[2J")
	fmt.Println("=== System Monitor ===")
	text := fmt.Sprintf("[CPU]\nLoad: %.2f%%  Avg: %.2f, %.2f, %.2f",
		percent[0],
		loadAvg.Load1,
		loadAvg.Load5,
		loadAvg.Load15)

	if showCores {
		perCPU, _ := cpu.Percent(time.Second, true)
		for i, p := range perCPU {
			text += fmt.Sprintf("\nCore %d: %s %.2f%%",
				i+1,
				textProgressBar(p, 20),
				p)
		}
	}
	return text
}

func getMemoryInfo() string {
	memStat, _ := mem.VirtualMemory()
	return fmt.Sprintf("[Memory]\nUsed: %.2f%% (%.2fG/%.2fG)\n%s",
		memStat.UsedPercent,
		float64(memStat.Used)/1024/1024/1024,
		float64(memStat.Total)/1024/1024/1024,
		textProgressBar(memStat.UsedPercent, 50))
}

func getNetworkInfo() string {
	interfaces, _ := net.IOCounters(true)
	now := time.Now()
	var sent, recv float64

	for _, iface := range interfaces {
		if iface.Name == "lo" {
			continue
		}

		stat, exists := netStats[iface.Name]
		if !exists {
			stat = &NetStat{
				LastBytesSent: iface.BytesSent,
				LastBytesRecv: iface.BytesRecv,
				LastTime:      now,
			}
			netStats[iface.Name] = stat
			continue
		}

		timeDiff := now.Sub(stat.LastTime).Seconds()
		sentDiff := float64(iface.BytesSent - stat.LastBytesSent)
		recvDiff := float64(iface.BytesRecv - stat.LastBytesRecv)

		stat.CurrentSentRate = sentDiff / timeDiff
		stat.CurrentRecvRate = recvDiff / timeDiff
		stat.LastBytesSent = iface.BytesSent
		stat.LastBytesRecv = iface.BytesRecv
		stat.LastTime = now

		sent += stat.CurrentSentRate
		recv += stat.CurrentRecvRate
	}

	return fmt.Sprintf("[Network]\nUpload: %s\nDownload: %s",
		formatSpeed(sent),
		formatSpeed(recv))
}

func getDiskIOInfo() string {
	ioStats, _ := disk.IOCounters()
	text := "[Disk I/O]"

	for name, io := range ioStats {
		last, exists := diskIOHistory[name]
		if exists {
			interval := float64(refreshInterval.Seconds())
			read := float64(io.ReadBytes-last.ReadBytes) / interval
			write := float64(io.WriteBytes-last.WriteBytes) / interval

			text += fmt.Sprintf("\n%-8s Read: %s  Write: %s",
				name,
				formatSpeed(read),
				formatSpeed(write))
		}
		diskIOHistory[name] = io
	}
	return text
}

func textProgressBar(percent float64, width int) string {
	filled := int(percent / 100 * float64(width))
	bar := ""
	for i := 0; i < width; i++ {
		if i < filled {
			bar += "■"
		} else {
			bar += "□"
		}
	}
	return bar
}