pstree.go

package tui

import (
	"fmt"
	"get-container/cmd/hytop/backend"
	"get-container/utils"
	"slices"
	"strings"

	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/charmbracelet/lipgloss/tree"
)

const (
	ModelPsTreeHeader = "    PID USER      DEVICE %CPU  %MEM  TIME      COMMAND"
)

var (
	ModelPsTreeStyle     = lipgloss.NewStyle().Foreground(lipgloss.Color("#000000")).Background(lipgloss.Color("#00fffb7a"))
	ModelPsTreeStyleKey  = lipgloss.NewStyle().Foreground(lipgloss.Color("#842cffff")).Background(lipgloss.Color("#00fffb7a")).Italic(true)
	ModelPsTreeStyleInfo = lipgloss.NewStyle().Foreground(lipgloss.Color("#e82e2eff")).Background(lipgloss.Color("#00fffb7a"))

	ModelPsTreeStyleL     = lipgloss.NewStyle().Foreground(lipgloss.Color("#656565b0")).Background(lipgloss.Color("#00fffb7a"))
	ModelPsTreeStyleKeyL  = lipgloss.NewStyle().Foreground(lipgloss.Color("#6a32b992")).Background(lipgloss.Color("#00fffb7a")).Italic(true)
	ModelPsTreeStyleInfoL = lipgloss.NewStyle().Foreground(lipgloss.Color("#e82e2e80")).Background(lipgloss.Color("#00fffb7a"))

	ModelPsFullHeightStyle = lipgloss.NewStyle()
)

type ModelPsTree struct {
	width, height int
	modelMsg      *ModelMsg
	actionMsg     *ActionMsg
	header        [2]string
	mapPidToDCU   map[int32]int // key为进程ID，value为dcu index

	treeWidth int // 树形图的最大宽度
}

func NewModelPsTree(width, height int) *ModelPsTree {
	return &ModelPsTree{
		width:       width,
		height:      height,
		mapPidToDCU: make(map[int32]int),
	}
}

func (m *ModelPsTree) Init() tea.Cmd {
	// 生成标题字符串
	sb := strings.Builder{}

	h1 := ModelPsTreeStyle.SetString(ModelPsTreeHeader).String()
	sb.WriteString(ModelPsTreeStyle.SetString("(Press ").String())
	sb.WriteString(ModelPsTreeStyleKey.SetString("^C").String())
	sb.WriteString(ModelPsTreeStyle.SetString("(").String())
	sb.WriteString(ModelPsTreeStyleInfo.SetString("INT").String())
	sb.WriteString(ModelPsTreeStyle.SetString(")/").String())
	sb.WriteString(ModelPsTreeStyleKey.SetString("T").String())
	sb.WriteString(ModelPsTreeStyle.SetString("(").String())
	sb.WriteString(ModelPsTreeStyleInfo.SetString("TERM").String())
	sb.WriteString(ModelPsTreeStyle.SetString(")/").String())
	sb.WriteString(ModelPsTreeStyleKey.SetString("K").String())
	sb.WriteString(ModelPsTreeStyle.SetString("(").String())
	sb.WriteString(ModelPsTreeStyleInfo.SetString("KILL").String())
	sb.WriteString(ModelPsTreeStyle.SetString(") to send signals)").String())
	h2 := sb.String()
	sb.Reset()
	w := m.width - lipgloss.Width(h1) - lipgloss.Width(h2)
	h := ModelPsTreeStyle.SetString(strings.Repeat(" ", w)).String()
	m.header[0] = h1 + h + h2

	h1 = ModelPsTreeStyleL.SetString(ModelPsTreeHeader).String()
	sb.WriteString(ModelPsTreeStyleL.SetString("(Press ").String())
	sb.WriteString(ModelPsTreeStyleKeyL.SetString("^C").String())
	sb.WriteString(ModelPsTreeStyleL.SetString("(").String())
	sb.WriteString(ModelPsTreeStyleInfoL.SetString("INT").String())
	sb.WriteString(ModelPsTreeStyleL.SetString(")/").String())
	sb.WriteString(ModelPsTreeStyleKeyL.SetString("T").String())
	sb.WriteString(ModelPsTreeStyleL.SetString("(").String())
	sb.WriteString(ModelPsTreeStyleInfoL.SetString("TERM").String())
	sb.WriteString(ModelPsTreeStyleL.SetString(")/").String())
	sb.WriteString(ModelPsTreeStyleKeyL.SetString("K").String())
	sb.WriteString(ModelPsTreeStyleL.SetString("(").String())
	sb.WriteString(ModelPsTreeStyleInfoL.SetString("KILL").String())
	sb.WriteString(ModelPsTreeStyleL.SetString(") to send signals)").String())
	h2 = sb.String()
	h = ModelPsTreeStyleL.SetString(strings.Repeat(" ", w)).String()
	m.header[1] = h1 + h + h2

	m.treeWidth = m.width - lipgloss.Width(ModelPsTreeHeader) + lipgloss.Width("COMMAND")

	ModelPsFullHeightStyle = ModelPsFullHeightStyle.Height(m.height).Width(m.width).MaxHeight(m.height).MaxWidth(m.width).Align(lipgloss.Center, lipgloss.Top)
	return nil
}

func (m *ModelPsTree) View() string {
	pids := make([]int32, 0, len(m.mapPidToDCU))
	for k := range m.mapPidToDCU {
		pids = append(pids, k)
	}
	slices.Sort(pids)
	pinfo, root := utils.GetPsTree(pids)
	if len(pinfo) > 0 && root != nil {
		darkMode := false
		if m.actionMsg.Action != nil && m.actionMsg.VM == VMTree {
			darkMode = true
		}
		tr := tree.New()
		tr.Root(root.Cmd)
		pidSort := genTree(tr, root, make([]int32, 0, 128))
		treeStr := lipgloss.NewStyle().MaxWidth(m.treeWidth).Render(tr.String())
		style := lipgloss.NewStyle().Inline(true)
		sb := strings.Builder{}
		for _, v := range pidSort {
			p, have := pinfo[v]
			if !have {
				sb.WriteByte('\n')
				continue
			}
			sb.WriteString(FormatStr(fmt.Sprintf("%d", p.Pid), 7, lipgloss.Right))
			sb.WriteByte(' ')
			sb.WriteString(style.MaxWidth(9).Width(9).Align(lipgloss.Left).Render(p.User))
			sb.WriteByte(' ')
			if dcu, have := m.mapPidToDCU[p.Pid]; have {
				sb.WriteString(style.MaxWidth(6).Width(6).Align(lipgloss.Left).Render(fmt.Sprintf("DCU%d", dcu)))
			} else {
				sb.WriteString(style.MaxWidth(6).Width(6).Align(lipgloss.Left).Render("Host"))
			}
			sb.WriteByte(' ')
			sb.WriteString(style.MaxWidth(5).Width(5).Align(lipgloss.Left).Render(fmt.Sprintf("%.1f", p.CPU)))
			sb.WriteByte(' ')
			sb.WriteString(style.MaxWidth(5).Width(5).Align(lipgloss.Left).Render(fmt.Sprintf("%.1f", p.Mem)))
			sb.WriteByte(' ')
			sb.WriteString(style.MaxWidth(9).Width(9).Align(lipgloss.Left).Render(backend.DurationStr(p.Time)))
			sb.WriteByte(' ')
			sb.WriteByte('\n')
		}
		infoStr := strings.Trim(sb.String(), "\n")
		if m.actionMsg != nil {
			if m.actionMsg.PidTreePids != nil {
				m.actionMsg.PidTreePids = m.actionMsg.PidTreePids[:0]
				m.actionMsg.PidTreePids = append(m.actionMsg.PidTreePids, pidSort...)
			} else {
				m.actionMsg.PidTreePids = make([]int32, 0, len(pidSort))
				m.actionMsg.PidTreePids = append(m.actionMsg.PidTreePids, pidSort...)
			}
		}
		sl := len(m.actionMsg.SelectPids)
		downStr := lipgloss.JoinHorizontal(lipgloss.Top, infoStr, treeStr)
		if darkMode {
			downStr = LowLeightStyle.Render(downStr)
			finalStr := lipgloss.JoinVertical(lipgloss.Left, m.header[1], downStr)
			return ModelPsFullHeightStyle.Render(finalStr)
		}
		lines := strings.Split(downStr, "\n")
		for index, pid := range pidSort {
			selected, pointed := false, false
			if sl > 0 {
				a, b := m.actionMsg.SelectPids[pid]
				selected = a && b
			}
			if m.actionMsg.PointPid != nil {
				pointed = *m.actionMsg.PointPid == pid
			}
			if selected && pointed {
				lines[index] = HeightSelectedStyle.Render(lines[index])
			} else if selected {
				lines[index] = SelectedStyle.Render(lines[index])
			} else if pointed {
				lines[index] = HeightLightStyle.Render(lines[index])
			}
		}
		downStr = strings.Join(lines, "\n")
		finalStr := lipgloss.JoinVertical(lipgloss.Left, m.header[0], downStr)
		return ModelPsFullHeightStyle.Render(finalStr)
	} else {
		return m.header[0] + "\n There is no DCU process running"
	}
}

func (m *ModelPsTree) Update(imsg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := imsg.(type) {
	case *ModelMsg:
		m.modelMsg = msg
		clear(m.mapPidToDCU)
		for k, v := range msg.DCUPidInfo {
			for _, val := range v {
				m.mapPidToDCU[val.Info.Pid] = k
			}
		}
	case *ActionMsg:
		m.actionMsg = msg
	}
	return m, nil
}

func genTree(t *tree.Tree, root *utils.ProcessInfo, ctx []int32) []int32 {
	ctx = append(ctx, root.Pid)
	pids := make([]int32, 0, len(root.Child))
	for k := range root.Child {
		pids = append(pids, k)
	}
	slices.Sort(pids)
	for _, v := range pids {
		c := root.Child[v]
		tt := tree.New().Root(c.Cmd)
		t.Child(tt)
		if len(c.Child) > 0 {
			ctx = genTree(tt, c, ctx)
		} else {
			ctx = append(ctx, c.Pid)
		}
	}
	return ctx
}