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

// Package model provides TUI data structures and state management
package model

import (
	"time"

	"github.com/charmbracelet/bubbles/list"
	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/textinput"
	"github.com/charmbracelet/bubbles/viewport"
	"github.com/jfraeys/fetch_ml/internal/domain"
)

// Re-export domain types for TUI use
// Task represents a task in the TUI
type Task = domain.Task

// TrackingConfig specifies experiment tracking tools
type TrackingConfig = domain.TrackingConfig

// MLflowTrackingConfig controls MLflow integration
type MLflowTrackingConfig = domain.MLflowTrackingConfig

// TensorBoardTrackingConfig controls TensorBoard integration
type TensorBoardTrackingConfig = domain.TensorBoardTrackingConfig

// WandbTrackingConfig controls Weights & Biases integration
type WandbTrackingConfig = domain.WandbTrackingConfig

// ViewMode represents the current view mode in the TUI
type ViewMode int

// ViewMode constants represent different TUI views
const (
	ViewModeJobs              ViewMode = iota // Jobs view mode
	ViewModeGPU                               // GPU status view mode
	ViewModeQueue                             // Queue status view mode
	ViewModeContainer                         // Container status view mode
	ViewModeSettings                          // Settings view mode
	ViewModeDatasets                          // Datasets view mode
	ViewModeExperiments                       // Experiments view mode
	ViewModeNarrative                         // Narrative/Outcome view mode
	ViewModeTeam                              // Team collaboration view mode
	ViewModeExperimentHistory                 // Experiment history view mode
	ViewModeConfig                            // Config view mode
	ViewModeLogs                              // Logs streaming view mode
)

// DatasetInfo represents dataset information in the TUI
type DatasetInfo struct {
	LastAccess time.Time `json:"last_access"`
	Name       string    `json:"name"`
	Location   string    `json:"location"`
	SizeBytes  int64     `json:"size_bytes"`
}

// State holds the application state
type State struct {
	JobList               list.Model
	LastRefresh           time.Time
	LastGPUUpdate         time.Time
	LastFrameTime         time.Time
	JobStats              map[JobStatus]int
	Status                string
	APIKey                string
	ErrorMsg              string
	Keys                  KeyMap
	QueuedTasks           []*Task
	Datasets              []DatasetInfo
	Jobs                  []Job
	Input                 textinput.Model
	APIKeyInput           textinput.Model
	GpuView               viewport.Model
	QueueView             viewport.Model
	LogsView              viewport.Model
	ConfigView            viewport.Model
	ExperimentHistoryView viewport.Model
	TeamView              viewport.Model
	SettingsView          viewport.Model
	DatasetView           viewport.Model
	ExperimentsView       viewport.Model
	NarrativeView         viewport.Model
	ContainerView         viewport.Model
	Spinner               spinner.Model
	SelectedJob           Job
	ActiveView            ViewMode
	RefreshRate           float64
	FrameCount            int
	Height                int
	Width                 int
	SettingsIndex         int
	ShowHelp              bool
	IsLoading             bool
	InputMode             bool
}

// InitialState creates the initial application state
func InitialState(apiKey string) State {
	items := []list.Item{}
	delegate := NewJobListDelegate()
	jobList := list.New(items, delegate, 0, 0)
	jobList.Title = "ML Jobs & Queue"
	jobList.SetShowStatusBar(true)
	jobList.SetFilteringEnabled(true)
	jobList.SetShowHelp(false)
	jobList.Styles.Title = JobListTitleStyle()

	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 = SpinnerStyle()

	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),
		NarrativeView:         viewport.New(0, 0),
		TeamView:              viewport.New(0, 0),
		ExperimentHistoryView: viewport.New(0, 0),
		ConfigView:            viewport.New(0, 0),
		LogsView:              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:                  DefaultKeys(),
	}
}

// LogMsg represents a log line from a job
type LogMsg struct {
	JobName string `json:"job_name"`
	Line    string `json:"line"`
	Time    string `json:"time"`
}

// JobUpdateMsg represents a real-time job status update via WebSocket
type JobUpdateMsg struct {
	JobName  string `json:"job_name"`
	Status   string `json:"status"`
	TaskID   string `json:"task_id"`
	Progress int    `json:"progress"`
}

// GPUUpdateMsg represents a real-time GPU status update via WebSocket
type GPUUpdateMsg struct {
	DeviceID    int   `json:"device_id"`
	Utilization int   `json:"utilization"`
	MemoryUsed  int64 `json:"memory_used"`
	MemoryTotal int64 `json:"memory_total"`
	Temperature int   `json:"temperature"`
}