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

// Package model provides TUI data structures and state management
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"
	"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
)

// JobStatus represents the status of a job
type JobStatus string

// JobStatus constants represent different job states
const (
	StatusPending  JobStatus = "pending"  // Job is pending
	StatusQueued   JobStatus = "queued"   // Job is queued
	StatusRunning  JobStatus = "running"  // Job is running
	StatusFinished JobStatus = "finished" // Job is finished
	StatusFailed   JobStatus = "failed"   // Job is failed
)

// Job represents a job in the TUI
type Job struct {
	Name     string
	Status   JobStatus
	TaskID   string
	Priority int64
}

// Title returns the job title for display
func (j Job) Title() string { return j.Name }

// Description returns a formatted description with status icon
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)
}

// FilterValue returns the value used for filtering
func (j Job) FilterValue() string { return j.Name }

// DatasetInfo represents dataset information in the TUI
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
}

// KeyMap defines key bindings for the TUI
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
}

// Keys contains the default key bindings for the TUI
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")),
}

// InitialState creates the initial application state
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,
	}
}