fetch_ml/cmd/tui/internal/model/state.go

package model

import (
	"fmt"
	"time"

	"github.com/charmbracelet/bubbles/key"
	"github.com/charmbracelet/bubbles/list"
	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/textinput"
	"github.com/charmbracelet/bubbles/viewport"
	"github.com/charmbracelet/lipgloss"
)

type ViewMode int

const (
	ViewModeJobs ViewMode = iota
	ViewModeGPU
	ViewModeQueue
	ViewModeContainer
	ViewModeSettings
	ViewModeDatasets
	ViewModeExperiments
)

type JobStatus string

const (
	StatusPending  JobStatus = "pending"
	StatusQueued   JobStatus = "queued"
	StatusRunning  JobStatus = "running"
	StatusFinished JobStatus = "finished"
	StatusFailed   JobStatus = "failed"
)

type Job struct {
	Name     string
	Status   JobStatus
	TaskID   string
	Priority int64
}

func (j Job) Title() string { return j.Name }
func (j Job) Description() string {
	icon := map[JobStatus]string{
		StatusPending:  "⏸",
		StatusQueued:   "⏳",
		StatusRunning:  "▶",
		StatusFinished: "✓",
		StatusFailed:   "✗",
	}[j.Status]
	pri := ""
	if j.Priority > 0 {
		pri = fmt.Sprintf(" [P%d]", j.Priority)
	}
	return fmt.Sprintf("%s %s%s", icon, j.Status, pri)
}
func (j Job) FilterValue() string { return j.Name }

type Task struct {
	ID        string            `json:"id"`
	JobName   string            `json:"job_name"`
	Args      string            `json:"args"`
	Status    string            `json:"status"`
	Priority  int64             `json:"priority"`
	CreatedAt time.Time         `json:"created_at"`
	StartedAt *time.Time        `json:"started_at,omitempty"`
	EndedAt   *time.Time        `json:"ended_at,omitempty"`
	Error     string            `json:"error,omitempty"`
	Metadata  map[string]string `json:"metadata,omitempty"`
}

type DatasetInfo struct {
	Name       string    `json:"name"`
	SizeBytes  int64     `json:"size_bytes"`
	Location   string    `json:"location"`
	LastAccess time.Time `json:"last_access"`
}

// State holds the application state
type State struct {
	Jobs            []Job
	QueuedTasks     []*Task
	Datasets        []DatasetInfo
	JobList         list.Model
	GpuView         viewport.Model
	ContainerView   viewport.Model
	QueueView       viewport.Model
	SettingsView    viewport.Model
	DatasetView     viewport.Model
	ExperimentsView viewport.Model
	Input           textinput.Model
	ApiKeyInput     textinput.Model
	Status          string
	ErrorMsg        string
	InputMode       bool
	Width           int
	Height          int
	ShowHelp        bool
	Spinner         spinner.Model
	ActiveView      ViewMode
	LastRefresh     time.Time
	IsLoading       bool
	JobStats        map[JobStatus]int
	ApiKey          string
	SettingsIndex   int
	Keys            KeyMap
}

type KeyMap struct {
	Refresh         key.Binding
	Trigger         key.Binding
	TriggerArgs     key.Binding
	ViewQueue       key.Binding
	ViewContainer   key.Binding
	ViewGPU         key.Binding
	ViewJobs        key.Binding
	ViewDatasets    key.Binding
	ViewExperiments key.Binding
	ViewSettings    key.Binding
	Cancel          key.Binding
	Delete          key.Binding
	MarkFailed      key.Binding
	RefreshGPU      key.Binding
	Help            key.Binding
	Quit            key.Binding
}

var Keys = KeyMap{
	Refresh:         key.NewBinding(key.WithKeys("r"), key.WithHelp("r", "refresh all")),
	Trigger:         key.NewBinding(key.WithKeys("t"), key.WithHelp("t", "queue job")),
	TriggerArgs:     key.NewBinding(key.WithKeys("a"), key.WithHelp("a", "queue w/ args")),
	ViewQueue:       key.NewBinding(key.WithKeys("v"), key.WithHelp("v", "view queue")),
	ViewContainer:   key.NewBinding(key.WithKeys("o"), key.WithHelp("o", "containers")),
	ViewGPU:         key.NewBinding(key.WithKeys("g"), key.WithHelp("g", "gpu status")),
	ViewJobs:        key.NewBinding(key.WithKeys("1"), key.WithHelp("1", "job list")),
	ViewDatasets:    key.NewBinding(key.WithKeys("2"), key.WithHelp("2", "datasets")),
	ViewExperiments: key.NewBinding(key.WithKeys("3"), key.WithHelp("3", "experiments")),
	Cancel:          key.NewBinding(key.WithKeys("c"), key.WithHelp("c", "cancel task")),
	Delete:          key.NewBinding(key.WithKeys("d"), key.WithHelp("d", "delete job")),
	MarkFailed:      key.NewBinding(key.WithKeys("f"), key.WithHelp("f", "mark failed")),
	RefreshGPU:      key.NewBinding(key.WithKeys("G"), key.WithHelp("G", "refresh GPU")),
	ViewSettings:    key.NewBinding(key.WithKeys("s"), key.WithHelp("s", "settings")),
	Help:            key.NewBinding(key.WithKeys("h", "?"), key.WithHelp("h/?", "toggle help")),
	Quit:            key.NewBinding(key.WithKeys("q", "ctrl+c"), key.WithHelp("q", "quit")),
}

func InitialState(apiKey string) State {
	items := []list.Item{}
	delegate := list.NewDefaultDelegate()
	delegate.Styles.SelectedTitle = delegate.Styles.SelectedTitle.
		Foreground(lipgloss.Color("170")).
		Bold(true)
	delegate.Styles.SelectedDesc = delegate.Styles.SelectedDesc.
		Foreground(lipgloss.Color("246"))

	jobList := list.New(items, delegate, 0, 0)
	jobList.Title = "ML Jobs & Queue"
	jobList.SetShowStatusBar(true)
	jobList.SetFilteringEnabled(true)
	jobList.SetShowHelp(false)
	// Styles will be set in View or here?
	// Keeping style initialization here as it's part of the model state setup
	jobList.Styles.Title = lipgloss.NewStyle().
		Bold(true).
		Foreground(lipgloss.AdaptiveColor{Light: "#2980b9", Dark: "#7aa2f7"}).
		Padding(0, 0, 1, 0)

	input := textinput.New()
	input.Placeholder = "Args: --epochs 100 --lr 0.001 --priority 5"
	input.Width = 60
	input.CharLimit = 200

	apiKeyInput := textinput.New()
	apiKeyInput.Placeholder = "Enter API key..."
	apiKeyInput.Width = 40
	apiKeyInput.CharLimit = 200

	s := spinner.New()
	s.Spinner = spinner.Dot
	s.Style = lipgloss.NewStyle().Foreground(lipgloss.AdaptiveColor{Light: "#2980b9", Dark: "#7aa2f7"})

	return State{
		JobList:         jobList,
		GpuView:         viewport.New(0, 0),
		ContainerView:   viewport.New(0, 0),
		QueueView:       viewport.New(0, 0),
		SettingsView:    viewport.New(0, 0),
		DatasetView:     viewport.New(0, 0),
		ExperimentsView: viewport.New(0, 0),
		Input:           input,
		ApiKeyInput:     apiKeyInput,
		Status:          "Connected",
		InputMode:       false,
		ShowHelp:        false,
		Spinner:         s,
		ActiveView:      ViewModeJobs,
		LastRefresh:     time.Now(),
		IsLoading:       false,
		JobStats:        make(map[JobStatus]int),
		ApiKey:          apiKey,
		SettingsIndex:   0,
		Keys:            Keys,
	}
}