main.go

package tui

import (
	"get-container/cmd/hytop/backend"
	"get-container/gpu"
	"get-container/utils"
	"slices"
	"time"

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

const (
	DCUTopVersion = "1.0.0"
)

var (
	HeightLightStyle    = lipgloss.NewStyle().Background(lipgloss.Color("#00fffb7a")) // 高亮风格
	SelectedStyle       = lipgloss.NewStyle().Foreground(lipgloss.Color("#ffb81fe3")) // 被选中的风格
	HeightSelectedStyle = lipgloss.NewStyle().Background(lipgloss.Color("#ffb81fe3")) // 既被选中，又高亮的风格

	LowLeightStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#442d2d"))
	NormalStyle    = lipgloss.NewStyle().Foreground(lipgloss.Color("#e6e6e6eb"))

	myBorder = lipgloss.Border{
		Top:          "═",
		TopLeft:      "╒",
		TopRight:     "╕",
		Bottom:       "═",
		BottomLeft:   "╘",
		BottomRight:  "╛",
		Left:         "│",
		Right:        "│",
		MiddleLeft:   "├",
		MiddleRight:  "┤",
		Middle:       "┼",
		MiddleTop:    "┬",
		MiddleBottom: "┴",
	}
)

// ModelMsg 模型信息，在父组件和各个子组件间共享信息
type ModelMsg struct {
	t          time.Time                        // 当前时间
	Version    *gpu.HYVersionInfo               // gpu版本相关信息
	DCUPidInfo map[int][]backend.DCUProcessInfo // 使用dcu的进程信息
	systemInfo *utils.SysInfo                   // 系统信息
	MyVersion  string                           // 本软件的版本
	DCUInfo    *backend.DCUInfoMap              // dcu相关信息
}

type TickMsg time.Time
type ProcessAction int
type ViewMode int

const (
	PAKill ProcessAction = 9
	PAInt  ProcessAction = 2
	PATerm ProcessAction = 15
	PANone ProcessAction = 0

	VMMain       ViewMode = 0
	VMTree       ViewMode = 1
	VMOneProcess ViewMode = 2
	VMProcessEnv ViewMode = 3
)

func (pa ProcessAction) String() string {
	switch pa {
	case PAKill:
		return "SIGKILL"
	case PAInt:
		return "SIGINT"
	case PATerm:
		return "SIGTERM"
	default:
		return "Cancel"
	}
}

// ActionMsg 动作消息
type ActionMsg struct {
	VM             ViewMode       // 所在视图，默认为0，即主视图
	SelectPids     map[int32]bool // 选择的pid进程
	Action         *ProcessAction // 对选择的pid的动作
	PidView        *int32         // 进程视图指标的pid号，为null表示没有进入进程指标视图
	PointPid       *int32         // 指针指向的pid，为null表示没有选择进程，在主视图和进程树视图起作用
	PidEnvView     *int32         // 进程环境变量视图的pid号，为null表示不进入进程环境变量视图
	PidTreePids    []int32        // pid tree视图下，进程id的顺序列表
	TargetDCUIndex *int           // PointPid使用的dcu index
}

// ModelMain tui主模型类
type ModelMain struct {
	width, height int // 终端尺寸
	Header        *ModelHeader
	DCUInfo       *ModelDCUInfo
	SysLoad       *ModelSysLoad
	ProcessInfo   *ModelProcessInfo
	modelMsg      *ModelMsg
	actionMsg     *ActionMsg
	dialog        *Dialog      // 对话框
	pstree        *ModelPsTree // 进程树视图
	processDetail *ModelProcessDetail
}

func NewModelMain(width, height int) ModelMain {
	result := ModelMain{}
	result.width = width
	result.height = height
	result.Header = NewModelHeader()
	result.DCUInfo = NewModelDCUInfo(width, height)
	result.SysLoad = NewModelSysLoad(width)
	result.ProcessInfo = NewModelProcessInfo(width)
	result.dialog = NewDialog(nil)
	result.pstree = NewModelPsTree(width, height)
	result.actionMsg = &ActionMsg{}
	return result
}

func tickCmd() tea.Cmd {
	return tea.Every(time.Second, func(t time.Time) tea.Msg {
		return TickMsg(t)
	})
}

func sendMsgCmd(modelMsg *ModelMsg) tea.Cmd {
	return func() tea.Msg {
		return *modelMsg
	}
}

// Init 初始化信息
func (m *ModelMain) Init() tea.Cmd {
	modelMsg := ModelMsg{}
	if err := initModelInfo(&modelMsg); err != nil {
		return tea.Quit
	}
	cmds := make([]tea.Cmd, 0)
	if c := m.Header.Init(); c != nil {
		cmds = append(cmds, c)
	}
	if c := m.DCUInfo.Init(); c != nil {
		cmds = append(cmds, c)
	}
	if c := m.SysLoad.Init(); c != nil {
		cmds = append(cmds, c)
	}
	if c := m.ProcessInfo.Init(); c != nil {
		cmds = append(cmds, c)
	}
	if c := m.dialog.Init(); c != nil {
		cmds = append(cmds, c)
	}
	if c := m.pstree.Init(); c != nil {
		cmds = append(cmds, c)
	}
	m.modelMsg = &modelMsg
	// 将调用初始化方法
	cmds = append(cmds, tickCmd(), sendMsgCmd(&modelMsg))
	return tea.Batch(cmds...)
}

func (m *ModelMain) Update(inputMsg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := inputMsg.(type) {
	case tea.KeyMsg: // 键盘事件
		switch msg.String() {
		case "ctrl+c", "q":
			return m, tea.Quit
		case "up":
			cmd := m.handleKeyUp()
			return m, cmd
		case "down":
			cmd := m.handleKeyDown()
			return m, cmd
		case "enter":
			cmd := m.handleKeyEnter()
			return m, cmd
		case "h":
			return m, tea.Quit
		case "left", "right":
			cmd := m.handleKeyLR(msg.String())
			return m, cmd
		case "k":
			cmd := m.handleKeyK()
			return m, cmd
		case "esc":
			cmd := m.handleKeyEsc()
			return m, cmd
		case " ":
			cmd := m.handleKeySpace()
			return m, cmd
		case "t":
			cmd := m.handleKeyT()
			return m, cmd
		}
	case TickMsg: // 定时事件
		updateModelInfo(m.modelMsg, time.Time(msg))
		header, _ := m.Header.Update(m.modelMsg)
		dcuInfo, _ := m.DCUInfo.Update(m.modelMsg)
		sysLoad, _ := m.SysLoad.Update(m.modelMsg)
		pidinfo, _ := m.ProcessInfo.Update(m.modelMsg)
		pstree, _ := m.pstree.Update(m.modelMsg)
		m.Header = header.(*ModelHeader)
		m.DCUInfo = dcuInfo.(*ModelDCUInfo)
		m.SysLoad = sysLoad.(*ModelSysLoad)
		m.ProcessInfo = pidinfo.(*ModelProcessInfo)
		m.pstree = pstree.(*ModelPsTree)
		if m.processDetail != nil {
			detail, _ := m.processDetail.Update(m.modelMsg)
			m.processDetail = detail.(*ModelProcessDetail)
		}
		return m, tickCmd()
	case ModelMsg: // 初始化
		header, _ := m.Header.Update(m.modelMsg)
		dcuInfo, _ := m.DCUInfo.Update(m.modelMsg)
		sysLoad, _ := m.SysLoad.Update(m.modelMsg)
		pidinfo, _ := m.ProcessInfo.Update(m.modelMsg)
		m.Header = header.(*ModelHeader)
		m.DCUInfo = dcuInfo.(*ModelDCUInfo)
		m.SysLoad = sysLoad.(*ModelSysLoad)
		m.ProcessInfo = pidinfo.(*ModelProcessInfo)
		return m, nil
	case tea.WindowSizeMsg: // 窗口尺寸变化
		m.width, m.height = msg.Width, msg.Height
		header, _ := m.Header.Update(msg)
		dcuInfo, _ := m.DCUInfo.Update(msg)
		sysLoad, _ := m.SysLoad.Update(msg)
		pidinfo, _ := m.ProcessInfo.Update(msg)
		m.Header = header.(*ModelHeader)
		m.DCUInfo = dcuInfo.(*ModelDCUInfo)
		m.SysLoad = sysLoad.(*ModelSysLoad)
		m.ProcessInfo = pidinfo.(*ModelProcessInfo)
		return m, nil
	}
	return m, nil
}

func (m *ModelMain) handleKeyEnter() tea.Cmd {
	if m.actionMsg.VM != VMMain || m.actionMsg.PointPid == nil {
		return nil
	}
	pid := *m.actionMsg.PointPid
	m.actionMsg.PidView = &pid
	m.processDetail = NewModelProcessDetail(m.width, m.height, pid)
	m.processDetail.Update(m.modelMsg)
	m.actionMsg.VM = VMOneProcess
	return nil
}

func (m *ModelMain) handleKeyT() tea.Cmd {
	// 在主视图下，没有选择任何进程时，可以进入进程树视图，且
	if m.actionMsg.VM == VMMain && len(m.actionMsg.SelectPids) == 0 {
		m.actionMsg.VM = VMTree
		pstree, _ := m.pstree.Update(m.actionMsg)
		m.pstree = pstree.(*ModelPsTree)
		return nil
	} else if m.actionMsg.VM == VMTree {
		if len(m.actionMsg.SelectPids) != 0 || m.actionMsg.PointPid != nil {
			action := PATerm
			m.actionMsg.Action = &action
			pstree, _ := m.pstree.Update(m.actionMsg)
			dialog, _ := m.dialog.Update(m.actionMsg)
			m.dialog = dialog.(*Dialog)
			m.pstree = pstree.(*ModelPsTree)
			return nil
		} else {
			return nil
		}
	}
	return nil
}

// handleKeyLR 处理左右键消息，左右键仅在对话框中有效
func (m *ModelMain) handleKeyLR(s string) tea.Cmd {
	if m.actionMsg == nil || m.actionMsg.Action == nil {
		return nil
	}
	switch s {
	case "right":
		switch *m.actionMsg.Action {
		case PAKill:
			*m.actionMsg.Action = PATerm
		case PATerm:
			*m.actionMsg.Action = PAInt
		case PAInt:
			*m.actionMsg.Action = PANone
		default:
			return nil
		}
	case "left":
		switch *m.actionMsg.Action {
		case PATerm:
			*m.actionMsg.Action = PAKill
		case PAInt:
			*m.actionMsg.Action = PATerm
		case PANone:
			*m.actionMsg.Action = PAInt
		default:
			return nil
		}
	default:
		return nil
	}
	dialog, _ := m.dialog.Update(m.actionMsg)
	m.dialog = dialog.(*Dialog)
	return nil
}

func (m *ModelMain) handleKeyK() tea.Cmd {
	if m.actionMsg.PointPid == nil && m.actionMsg.SelectPids == nil {
		return nil
	}
	if m.actionMsg.Action != nil {
		return nil
	}
	var pa ProcessAction = PANone
	if m.actionMsg.Action == nil {
		m.actionMsg.Action = &pa
	}
	switch m.actionMsg.VM {
	case VMMain:
		header, _ := m.Header.Update(m.actionMsg)
		dcuInfo, _ := m.DCUInfo.Update(m.actionMsg)
		sysLoad, _ := m.SysLoad.Update(m.actionMsg)
		pidinfo, _ := m.ProcessInfo.Update(m.actionMsg)
		m.Header = header.(*ModelHeader)
		m.DCUInfo = dcuInfo.(*ModelDCUInfo)
		m.SysLoad = sysLoad.(*ModelSysLoad)
		m.ProcessInfo = pidinfo.(*ModelProcessInfo)
		dialog, _ := m.dialog.Update(m.actionMsg)
		m.dialog = dialog.(*Dialog)
		return nil
	case VMTree:
		pstree, _ := m.pstree.Update(m.actionMsg)
		dialog, _ := m.dialog.Update(m.actionMsg)
		m.dialog = dialog.(*Dialog)
		m.pstree = pstree.(*ModelPsTree)
		return nil
	default:
		return nil
	}
}

// handleKeyEsc 处理Esc键，esc键在：
// VMMain视图下，取消所有选择的进程、退出单个进程模式或对话框返回
// VMTree视图下，取消所有选择的进程、退出单个进程模式或对话框返回，返回到VMMain视图
func (m *ModelMain) handleKeyEsc() tea.Cmd {
	switch m.actionMsg.VM {
	case VMMain:
		if m.actionMsg.Action != nil {
			m.actionMsg.Action = nil
		} else if m.actionMsg.SelectPids != nil {
			m.actionMsg.SelectPids = nil
		} else if m.actionMsg.PointPid != nil || m.actionMsg.TargetDCUIndex != nil {
			m.actionMsg.PointPid = nil
			m.actionMsg.TargetDCUIndex = nil
		}
		header, _ := m.Header.Update(m.actionMsg)
		dcuInfo, _ := m.DCUInfo.Update(m.actionMsg)
		sysLoad, _ := m.SysLoad.Update(m.actionMsg)
		pidinfo, _ := m.ProcessInfo.Update(m.actionMsg)
		m.Header = header.(*ModelHeader)
		m.DCUInfo = dcuInfo.(*ModelDCUInfo)
		m.SysLoad = sysLoad.(*ModelSysLoad)
		m.ProcessInfo = pidinfo.(*ModelProcessInfo)
		return nil
	case VMTree:
		if m.actionMsg.Action != nil {
			m.actionMsg.Action = nil
		} else if m.actionMsg.SelectPids != nil {
			m.actionMsg.SelectPids = nil
		} else if m.actionMsg.PointPid != nil {
			m.actionMsg.PointPid = nil
		} else {
			m.actionMsg.VM = VMMain
		}
		pstree, _ := m.pstree.Update(m.actionMsg)
		m.pstree = pstree.(*ModelPsTree)
		return nil
	case VMOneProcess:
		m.actionMsg.VM = VMMain
		m.processDetail = nil
		return nil
	default:
		return nil
	}
}

// handleKeySpace 处理空格键动作，空格键仅用于选择或取消进程
func (m *ModelMain) handleKeySpace() tea.Cmd {
	if len(m.modelMsg.DCUPidInfo) == 0 {
		return nil
	}
	if m.actionMsg.PointPid == nil || m.actionMsg.Action != nil {
		return nil
	}
	if m.actionMsg.VM != VMMain && m.actionMsg.VM != VMTree {
		return nil
	}
	if m.actionMsg.SelectPids == nil {
		m.actionMsg.SelectPids = make(map[int32]bool)
	}
	_, have := m.actionMsg.SelectPids[*m.actionMsg.PointPid]
	if have {
		delete(m.actionMsg.SelectPids, *m.actionMsg.PointPid)
	} else {
		m.actionMsg.SelectPids[*m.actionMsg.PointPid] = true
	}
	switch m.actionMsg.VM {
	case VMMain:
		m1, cmd1 := m.ProcessInfo.Update(m.actionMsg)
		m.ProcessInfo = m1.(*ModelProcessInfo)
		return cmd1
	case VMTree:
		// todo
		m1, cmd1 := m.pstree.Update(m.actionMsg)
		m.pstree = m1.(*ModelPsTree)
		return cmd1
	default:
		return nil
	}
}

func (m *ModelMain) View() string {
	switch m.actionMsg.VM {
	case VMMain:
		if m.actionMsg.Action != nil {
			up := m.dialog.View()
			down := m.Header.View() + m.DCUInfo.View() + m.SysLoad.View() + m.ProcessInfo.View() + "\n"
			return StackPosition(up, down, lipgloss.Center, lipgloss.Center)
		}
		return m.Header.View() + m.DCUInfo.View() + m.SysLoad.View() + m.ProcessInfo.View() + "\n"
	case VMTree:
		if m.actionMsg.Action != nil {
			up := m.dialog.View()
			down := m.pstree.View()
			return StackPosition(up, down, lipgloss.Center, lipgloss.Center)
		}
		return m.pstree.View()
	case VMOneProcess:
		return m.processDetail.View()
	default:
		return ""
	}
}

func initModelInfo(model *ModelMsg) error {
	model.t = time.Now()
	model.MyVersion = DCUTopVersion
	if ver, err := gpu.GetHYVersionInfo(); err != nil {
		return err
	} else {
		model.Version = ver
	}
	model.DCUInfo = backend.DCUSInfoMap
	model.DCUPidInfo = backend.DCUSInfoMap.GetDCUProcessInfo2()
	if sinfo, err := utils.GetSysInfo(); err != nil {
		return err
	} else {
		model.systemInfo = sinfo
	}
	if err := model.DCUInfo.UpdateQuickInfo(); err != nil {
		return err
	}
	if err := model.DCUInfo.UpdateSlowInfo(); err != nil {
		return err
	}
	return nil
}

// updateModelInfo 更新模型信息
func updateModelInfo(modelMsg *ModelMsg, t time.Time) error {
	modelMsg.t = t
	if sinfo, err := utils.GetSysInfo(); err != nil {
		return err
	} else {
		modelMsg.systemInfo = sinfo
	}
	modelMsg.DCUPidInfo = modelMsg.DCUInfo.GetDCUProcessInfo2()
	if err := modelMsg.DCUInfo.UpdateQuickInfo(); err != nil {
		return err
	}
	return nil
}

func (m *ModelMain) handleKeyUp() tea.Cmd {
	if len(m.modelMsg.DCUPidInfo) == 0 {
		return nil
	}
	if m.actionMsg.Action != nil {
		return nil
	}
	switch m.actionMsg.VM {
	case VMMain:
		processes := make([]backend.DCUProcessInfo, 0, 16)
		keys := make([]int, 0, 16)
		for k := range m.modelMsg.DCUPidInfo {
			keys = append(keys, k)
		}
		slices.Sort(keys)
		for _, index := range keys {
			p, have := m.modelMsg.DCUPidInfo[index]
			if have && len(p) > 0 {
				processes = append(processes, p...)
			}
		}
		processNum := len(processes)
		if processNum == 0 {
			return nil
		}
		index := 0
		if m.actionMsg.PointPid == nil {
			// 获取列表中的最后一个进程的pid
			index = processNum - 1
		} else {
			var targetIndex = -1
			for l := processNum - 1; l >= 0; l-- {
				if processes[l].Info.Pid == *m.actionMsg.PointPid {
					targetIndex = l - 1
					break
				}
			}
			if targetIndex == -1 {
				index = processNum - 1
			} else {
				index = targetIndex
			}
		}
		pid := processes[index].Info.Pid
		m.actionMsg.PointPid = &pid
		idx := processes[index].DCU
		m.actionMsg.TargetDCUIndex = &idx
		m1, cmd1 := m.ProcessInfo.Update(m.actionMsg)
		m2, cmd2 := m.DCUInfo.Update(m.actionMsg)
		m3, cmd3 := m.SysLoad.Update(m.actionMsg)
		m.ProcessInfo = m1.(*ModelProcessInfo)
		m.DCUInfo = m2.(*ModelDCUInfo)
		m.SysLoad = m3.(*ModelSysLoad)
		return tea.Batch(cmd1, cmd2, cmd3)
	case VMTree:
		if len(m.actionMsg.PidTreePids) == 0 {
			return nil
		}
		if m.actionMsg.PointPid == nil {
			pid := m.actionMsg.PidTreePids[len(m.actionMsg.PidTreePids)-1]
			m.actionMsg.PointPid = &pid
		} else {
			idx := slices.Index(m.actionMsg.PidTreePids, *m.actionMsg.PointPid)
			if idx == -1 || idx == 0 {
				*m.actionMsg.PointPid = m.actionMsg.PidTreePids[len(m.actionMsg.PidTreePids)-1]
			} else {
				*m.actionMsg.PointPid = m.actionMsg.PidTreePids[idx-1]
			}
		}
		pstree, _ := m.pstree.Update(m.actionMsg)
		m.pstree = pstree.(*ModelPsTree)
		return nil
	default:
		return nil
	}

}

func (m *ModelMain) handleKeyDown() tea.Cmd {
	if len(m.modelMsg.DCUPidInfo) == 0 {
		return nil
	}
	if m.actionMsg.Action != nil {
		return nil
	}
	switch m.actionMsg.VM {
	case VMMain:
		processes := make([]backend.DCUProcessInfo, 0, 16)
		keys := make([]int, 0, 16)
		for key := range m.modelMsg.DCUPidInfo {
			keys = append(keys, key)
		}
		slices.Sort(keys)
		for _, index := range keys {
			p, have := m.modelMsg.DCUPidInfo[index]
			if have && len(p) > 0 {
				processes = append(processes, p...)
			}
		}
		processNum := len(processes)
		if processNum == 0 {
			return nil
		}
		index := 0
		if m.actionMsg.PointPid == nil {
			index = 0
		} else {
			var targetIndex = -1
			for l := 0; l < processNum-1; l++ {
				if processes[l].Info.Pid == *m.actionMsg.PointPid {
					targetIndex = l + 1
					if targetIndex >= processNum {
						targetIndex = 0
					}
					break
				}
			}
			if targetIndex == -1 {
				index = 0
			} else {
				index = targetIndex

			}
		}
		pid := processes[index].Info.Pid
		m.actionMsg.PointPid = &pid
		idx := processes[index].DCU
		m.actionMsg.TargetDCUIndex = &idx
		m1, cmd1 := m.ProcessInfo.Update(m.actionMsg)
		m2, cmd2 := m.DCUInfo.Update(m.actionMsg)
		m3, cmd3 := m.SysLoad.Update(m.actionMsg)
		m.ProcessInfo = m1.(*ModelProcessInfo)
		m.DCUInfo = m2.(*ModelDCUInfo)
		m.SysLoad = m3.(*ModelSysLoad)
		return tea.Batch(cmd1, cmd2, cmd3)
	case VMTree:
		if len(m.actionMsg.PidTreePids) == 0 {
			return nil
		}
		if m.actionMsg.PointPid == nil {
			pid := m.actionMsg.PidTreePids[0]
			m.actionMsg.PointPid = &pid
		} else {
			idx := slices.Index(m.actionMsg.PidTreePids, *m.actionMsg.PointPid)
			if idx == -1 || idx == len(m.actionMsg.PidTreePids)-1 {
				*m.actionMsg.PointPid = m.actionMsg.PidTreePids[0]
			} else {
				*m.actionMsg.PointPid = m.actionMsg.PidTreePids[idx+1]
			}
		}
		pstree, _ := m.pstree.Update(m.actionMsg)
		m.pstree = pstree.(*ModelPsTree)
		return nil
	default:
		return nil
	}

}