fetch_ml/cmd/tui/internal/controller/controller.go

package controller

import (
	"fmt"
	"strings"
	"time"

	"github.com/charmbracelet/bubbles/key"
	"github.com/charmbracelet/bubbles/list"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/jfraeys/fetch_ml/cmd/tui/internal/config"
	"github.com/jfraeys/fetch_ml/cmd/tui/internal/model"
	"github.com/jfraeys/fetch_ml/cmd/tui/internal/services"
	"github.com/jfraeys/fetch_ml/internal/logging"
)

// Controller handles all business logic and state updates
type Controller struct {
	config    *config.Config
	server    *services.MLServer
	taskQueue *services.TaskQueue
	logger    *logging.Logger
	wsClient  *services.WebSocketClient
}

func (c *Controller) handleKeyMsg(msg tea.KeyMsg, m model.State) (model.State, tea.Cmd) {
	if m.InputMode {
		return c.handleInputModeKey(msg, m)
	}

	if m.ActiveView == model.ViewModeSettings {
		return c.handleSettingsKeys(msg, m)
	}

	if key.Matches(msg, m.Keys.Quit) {
		return m, tea.Quit
	}

	cmds := c.handleGlobalKeys(msg, &m)
	return c.finalizeUpdate(msg, m, cmds...)
}

func (c *Controller) handleInputModeKey(msg tea.KeyMsg, m model.State) (model.State, tea.Cmd) {
	switch msg.String() {
	case "enter":
		args := m.Input.Value()
		m.Input.SetValue("")
		m.InputMode = false
		if job := getSelectedJob(m); job != nil {
			return m, c.queueJob(job.Name, args)
		}
		return m, nil
	case "esc":
		m.InputMode = false
		m.Input.SetValue("")
		return m, nil
	default:
		var cmd tea.Cmd
		m.Input, cmd = m.Input.Update(msg)
		return m, cmd
	}
}

func (c *Controller) handleSettingsKeys(msg tea.KeyMsg, m model.State) (model.State, tea.Cmd) {
	cmds := c.navigateSettings(msg, &m)
	if m.SettingsIndex == 1 {
		var inputCmd tea.Cmd
		m.APIKeyInput, inputCmd = m.APIKeyInput.Update(msg)
		cmds = append(cmds, inputCmd, c.updateSettingsContent(m))
	}
	return m, tea.Batch(cmds...)
}

func (c *Controller) navigateSettings(msg tea.KeyMsg, m *model.State) []tea.Cmd {
	var cmds []tea.Cmd
	switch msg.String() {
	case "up", "k":
		if m.SettingsIndex > 1 {
			m.SettingsIndex--
			cmds = append(cmds, c.updateSettingsContent(*m))
			c.toggleAPIKeyInputFocus(m)
		}
	case "down", "j":
		if m.SettingsIndex < 2 {
			m.SettingsIndex++
			cmds = append(cmds, c.updateSettingsContent(*m))
			c.toggleAPIKeyInputFocus(m)
		}
	case "enter":
		if cmd := c.handleSettingsAction(m); cmd != nil {
			cmds = append(cmds, cmd)
		}
	case "esc":
		m.ActiveView = model.ViewModeJobs
		m.APIKeyInput.Blur()
	}
	return cmds
}

func (c *Controller) toggleAPIKeyInputFocus(m *model.State) {
	if m.SettingsIndex == 1 {
		m.APIKeyInput.Focus()
	} else {
		m.APIKeyInput.Blur()
	}
}

func (c *Controller) handleGlobalKeys(msg tea.KeyMsg, m *model.State) []tea.Cmd {
	var cmds []tea.Cmd

	switch {
	case key.Matches(msg, m.Keys.Refresh):
		m.IsLoading = true
		m.Status = "Refreshing all data..."
		m.LastRefresh = time.Now()
		cmds = append(cmds, c.loadAllData())
	case key.Matches(msg, m.Keys.RefreshGPU):
		m.Status = "Refreshing GPU status..."
		cmds = append(cmds, c.loadGPU())
	case key.Matches(msg, m.Keys.Trigger):
		if job := getSelectedJob(*m); job != nil {
			cmds = append(cmds, c.queueJob(job.Name, ""))
		}
	case key.Matches(msg, m.Keys.TriggerArgs):
		if job := getSelectedJob(*m); job != nil {
			m.InputMode = true
			m.Input.Focus()
		}
	case key.Matches(msg, m.Keys.ViewQueue):
		m.ActiveView = model.ViewModeQueue
		cmds = append(cmds, c.showQueue(*m))
	case key.Matches(msg, m.Keys.ViewContainer):
		m.ActiveView = model.ViewModeContainer
		cmds = append(cmds, c.loadContainer())
	case key.Matches(msg, m.Keys.ViewGPU):
		m.ActiveView = model.ViewModeGPU
		cmds = append(cmds, c.loadGPU())
	case key.Matches(msg, m.Keys.ViewJobs):
		m.ActiveView = model.ViewModeJobs
	case key.Matches(msg, m.Keys.ViewSettings):
		m.ActiveView = model.ViewModeSettings
		m.SettingsIndex = 1
		m.APIKeyInput.Focus()
		cmds = append(cmds, c.updateSettingsContent(*m))
	case key.Matches(msg, m.Keys.ViewExperiments):
		m.ActiveView = model.ViewModeExperiments
		cmds = append(cmds, c.loadExperiments())
	case key.Matches(msg, m.Keys.ViewNarrative):
		m.ActiveView = model.ViewModeNarrative
		if job := getSelectedJob(*m); job != nil {
			m.SelectedJob = *job
		}
	case key.Matches(msg, m.Keys.ViewTeam):
		m.ActiveView = model.ViewModeTeam
	case key.Matches(msg, m.Keys.ViewExperimentHistory):
		m.ActiveView = model.ViewModeExperimentHistory
		cmds = append(cmds, c.loadExperimentHistory())
	case key.Matches(msg, m.Keys.ViewConfig):
		m.ActiveView = model.ViewModeConfig
		cmds = append(cmds, c.loadConfig())
	case key.Matches(msg, m.Keys.ViewLogs):
		m.ActiveView = model.ViewModeLogs
		if job := getSelectedJob(*m); job != nil {
			cmds = append(cmds, c.loadLogs(job.Name))
		}
	case key.Matches(msg, m.Keys.ViewExport):
		if job := getSelectedJob(*m); job != nil {
			cmds = append(cmds, c.exportJob(job.Name))
		}
	case key.Matches(msg, m.Keys.FilterTeam):
		m.InputMode = true
		m.Input.SetValue("@")
		m.Input.Focus()
		m.Status = "Filter by team member: @alice, @bob, @team-ml"
	case key.Matches(msg, m.Keys.Cancel):
		if job := getSelectedJob(*m); job != nil && job.TaskID != "" {
			cmds = append(cmds, c.cancelTask(job.TaskID))
		}
	case key.Matches(msg, m.Keys.Delete):
		if job := getSelectedJob(*m); job != nil && job.Status == model.StatusPending {
			cmds = append(cmds, c.deleteJob(job.Name))
		}
	case key.Matches(msg, m.Keys.MarkFailed):
		if job := getSelectedJob(*m); job != nil && job.Status == model.StatusRunning {
			cmds = append(cmds, c.markFailed(job.Name))
		}
	case key.Matches(msg, m.Keys.Help):
		m.ShowHelp = !m.ShowHelp
	}

	return cmds
}

func (c *Controller) applyWindowSize(msg tea.WindowSizeMsg, m model.State) model.State {
	m.Width = msg.Width
	m.Height = msg.Height

	h, v := 4, 2
	listHeight := msg.Height - v - 8
	m.JobList.SetSize(msg.Width/3-h, listHeight)

	panelWidth := msg.Width*2/3 - h - 2
	panelHeight := (listHeight - 6) / 3

	m.GpuView.Width = panelWidth
	m.GpuView.Height = panelHeight
	m.ContainerView.Width = panelWidth
	m.ContainerView.Height = panelHeight
	m.QueueView.Width = panelWidth
	m.QueueView.Height = listHeight - 4
	m.SettingsView.Width = panelWidth
	m.SettingsView.Height = listHeight - 4
	m.NarrativeView.Width = panelWidth
	m.NarrativeView.Height = listHeight - 4
	m.TeamView.Width = panelWidth
	m.TeamView.Height = listHeight - 4
	m.ExperimentsView.Width = panelWidth
	m.ExperimentsView.Height = listHeight - 4
	m.ExperimentHistoryView.Width = panelWidth
	m.ExperimentHistoryView.Height = listHeight - 4
	m.ConfigView.Width = panelWidth
	m.ConfigView.Height = listHeight - 4
	m.LogsView.Width = panelWidth
	m.LogsView.Height = listHeight - 4

	return m
}

func (c *Controller) handleJobsLoadedMsg(msg model.JobsLoadedMsg, m model.State) (model.State, tea.Cmd) {
	m.Jobs = []model.Job(msg)
	calculateJobStats(&m)

	items := make([]list.Item, len(m.Jobs))
	for i, job := range m.Jobs {
		items[i] = job
	}

	setItemsCmd := m.JobList.SetItems(items)
	m.Status = formatStatus(m)
	m.IsLoading = false
	return c.finalizeUpdate(msg, m, setItemsCmd)
}

func (c *Controller) handleTasksLoadedMsg(
	msg model.TasksLoadedMsg,
	m model.State,
) (model.State, tea.Cmd) {
	m.QueuedTasks = []*model.Task(msg)
	m.Status = formatStatus(m)
	return c.finalizeUpdate(msg, m)
}

func (c *Controller) handleGPUContent(msg model.GpuLoadedMsg, m model.State) (model.State, tea.Cmd) {
	m.GpuView.SetContent(string(msg))
	m.GpuView.GotoTop()
	return c.finalizeUpdate(msg, m)
}

func (c *Controller) handleContainerContent(
	msg model.ContainerLoadedMsg,
	m model.State,
) (model.State, tea.Cmd) {
	m.ContainerView.SetContent(string(msg))
	m.ContainerView.GotoTop()
	return c.finalizeUpdate(msg, m)
}

func (c *Controller) handleQueueContent(msg model.QueueLoadedMsg, m model.State) (model.State, tea.Cmd) {
	m.QueueView.SetContent(string(msg))
	m.QueueView.GotoTop()
	return c.finalizeUpdate(msg, m)
}

func (c *Controller) handleStatusMsg(msg model.StatusMsg, m model.State) (model.State, tea.Cmd) {
	if msg.Level == "error" {
		m.ErrorMsg = msg.Text
		m.Status = "Error occurred - check status"
	} else {
		m.ErrorMsg = ""
		m.Status = msg.Text
	}
	return c.finalizeUpdate(msg, m)
}

func (c *Controller) handleTickMsg(msg model.TickMsg, m model.State) (model.State, tea.Cmd) {
	var cmds []tea.Cmd

	// Calculate actual refresh rate
	now := time.Now()
	if !m.LastFrameTime.IsZero() {
		elapsed := now.Sub(m.LastFrameTime).Milliseconds()
		if elapsed > 0 {
			// Smooth the rate with simple averaging
			m.RefreshRate = (m.RefreshRate*float64(m.FrameCount) + float64(elapsed)) / float64(m.FrameCount+1)
			m.FrameCount++
			if m.FrameCount > 100 {
				m.FrameCount = 1
				m.RefreshRate = float64(elapsed)
			}
		}
	}
	m.LastFrameTime = now

	// 500ms refresh target for real-time updates
	if time.Since(m.LastRefresh) > 500*time.Millisecond && !m.IsLoading {
		m.LastRefresh = time.Now()
		cmds = append(cmds, c.loadAllData())
	}
	cmds = append(cmds, tickCmd())
	return c.finalizeUpdate(msg, m, cmds...)
}

func (c *Controller) finalizeUpdate(
	msg tea.Msg,
	m model.State,
	extraCmds ...tea.Cmd,
) (model.State, tea.Cmd) {
	cmds := append([]tea.Cmd{}, extraCmds...)

	var cmd tea.Cmd
	m.JobList, cmd = m.JobList.Update(msg)
	cmds = append(cmds, cmd)

	m.GpuView, cmd = m.GpuView.Update(msg)
	cmds = append(cmds, cmd)

	m.ContainerView, cmd = m.ContainerView.Update(msg)
	cmds = append(cmds, cmd)

	m.QueueView, cmd = m.QueueView.Update(msg)
	cmds = append(cmds, cmd)

	m.ExperimentsView, cmd = m.ExperimentsView.Update(msg)
	cmds = append(cmds, cmd)

	var spinCmd tea.Cmd
	m.Spinner, spinCmd = m.Spinner.Update(msg)
	cmds = append(cmds, spinCmd)

	return m, tea.Batch(cmds...)
}

// New creates a new Controller instance
func New(
	cfg *config.Config,
	srv *services.MLServer,
	tq *services.TaskQueue,
	logger *logging.Logger,
) *Controller {
	// Create WebSocket client for real-time updates
	wsClient := services.NewWebSocketClient(cfg.ServerURL, "", logger)

	return &Controller{
		config:    cfg,
		server:    srv,
		taskQueue: tq,
		logger:    logger,
		wsClient:  wsClient,
	}
}

// Init initializes the TUI and returns initial commands
func (c *Controller) Init() tea.Cmd {
	// Connect WebSocket for real-time updates
	if err := c.wsClient.Connect(); err != nil {
		c.logger.Error("WebSocket connection failed", "error", err)
	}

	return tea.Batch(
		tea.SetWindowTitle("FetchML"),
		c.loadAllData(),
		tickCmd(),
	)
}

func (c *Controller) Update(msg tea.Msg, m model.State) (model.State, tea.Cmd) {
	switch typed := msg.(type) {
	case tea.KeyMsg:
		return c.handleKeyMsg(typed, m)
	case tea.WindowSizeMsg:
		// Only apply window size on first render, then keep constant
		if m.Width == 0 && m.Height == 0 {
			updated := c.applyWindowSize(typed, m)
			return c.finalizeUpdate(msg, updated)
		}
		return c.finalizeUpdate(msg, m)
	case model.JobsLoadedMsg:
		return c.handleJobsLoadedMsg(typed, m)
	case model.TasksLoadedMsg:
		return c.handleTasksLoadedMsg(typed, m)
	case model.GpuLoadedMsg:
		return c.handleGPUContent(typed, m)
	case model.ContainerLoadedMsg:
		return c.handleContainerContent(typed, m)
	case model.QueueLoadedMsg:
		return c.handleQueueContent(typed, m)
	case model.DatasetsLoadedMsg:
		// Format datasets into view content
		var content strings.Builder
		content.WriteString("Available Datasets\n")
		content.WriteString(strings.Repeat("═", 50) + "\n\n")
		if len(typed) == 0 {
			content.WriteString("📭 No datasets found\n\n")
			content.WriteString("Datasets will appear here when available\n")
			content.WriteString("in the data directory.")
		} else {
			for i, ds := range typed {
				content.WriteString(fmt.Sprintf("%d. 📁 %s\n", i+1, ds.Name))
				content.WriteString(fmt.Sprintf("   Location: %s\n", ds.Location))
				content.WriteString(fmt.Sprintf("   Size: %d bytes\n", ds.SizeBytes))
				content.WriteString(fmt.Sprintf("   Last Access: %s\n\n", ds.LastAccess.Format("2006-01-02 15:04")))
			}
		}
		m.DatasetView.SetContent(content.String())
		m.DatasetView.GotoTop()
		return c.finalizeUpdate(msg, m)
	case model.SettingsContentMsg:
		m.SettingsView.SetContent(string(typed))
		return c.finalizeUpdate(msg, m)
	case ExperimentsLoadedMsg:
		m.ExperimentsView.SetContent(string(typed))
		m.ExperimentsView.GotoTop()
		return c.finalizeUpdate(msg, m)
	case ExperimentHistoryLoadedMsg:
		m.ExperimentHistoryView.SetContent(string(typed))
		m.ExperimentHistoryView.GotoTop()
		return c.finalizeUpdate(msg, m)
	case ConfigLoadedMsg:
		m.ConfigView.SetContent(string(typed))
		m.ConfigView.GotoTop()
		return c.finalizeUpdate(msg, m)
	case LogsLoadedMsg:
		m.LogsView.SetContent(string(typed))
		m.LogsView.GotoTop()
		return c.finalizeUpdate(msg, m)
	case model.SettingsUpdateMsg:
		return c.finalizeUpdate(msg, m)
	case model.StatusMsg:
		return c.handleStatusMsg(typed, m)
	case model.TickMsg:
		return c.handleTickMsg(typed, m)
	case model.JobUpdateMsg:
		// Handle real-time job status updates from WebSocket
		m.Status = fmt.Sprintf("Job %s: %s", typed.JobName, typed.Status)
		// Refresh job list to show updated status
		return m, c.loadAllData()
	case model.GPUUpdateMsg:
		// Throttle GPU updates to 1/second (humans can't perceive faster)
		if time.Since(m.LastGPUUpdate) > 1*time.Second {
			m.LastGPUUpdate = time.Now()
			return c.finalizeUpdate(msg, m)
		}
		return m, nil
	default:
		return c.finalizeUpdate(msg, m)
	}
}

// ExperimentsLoadedMsg is sent when experiments are loaded
type ExperimentsLoadedMsg string

// ExperimentHistoryLoadedMsg is sent when experiment history is loaded
type ExperimentHistoryLoadedMsg string

// ConfigLoadedMsg is sent when config is loaded
type ConfigLoadedMsg string

func (c *Controller) loadExperiments() tea.Cmd {
	return func() tea.Msg {
		commitIDs, err := c.taskQueue.ListExperiments()
		if err != nil {
			return model.StatusMsg{Level: "error", Text: fmt.Sprintf("Failed to list experiments: %v", err)}
		}

		if len(commitIDs) == 0 {
			return ExperimentsLoadedMsg("Experiments:\n\nNo experiments found.")
		}

		var output string
		output += "Experiments:\n\n"

		for _, commitID := range commitIDs {
			details, err := c.taskQueue.GetExperimentDetails(commitID)
			if err != nil {
				output += fmt.Sprintf("Error loading %s: %v\n\n", commitID, err)
				continue
			}
			output += details + "\n----------------------------------------\n\n"
		}

		return ExperimentsLoadedMsg(output)
	}
}

func (c *Controller) loadExperimentHistory() tea.Cmd {
	return func() tea.Msg {
		// Placeholder - will show experiment history with annotations
		return ExperimentHistoryLoadedMsg("Experiment History & Annotations\n\n" +
			"This view will show:\n" +
			"- Previous experiment runs\n" +
			"- Annotations and notes\n" +
			"- Config snapshots\n" +
			"- Side-by-side comparisons\n\n" +
			"(Requires API: GET /api/experiments/:id/history)")
	}
}

func (c *Controller) loadConfig() tea.Cmd {
	return func() tea.Msg {
		// Build config diff showing changes from defaults
		var output strings.Builder
		output.WriteString("⚙️  Config View (Read-Only)\n\n")

		output.WriteString("┌─ Changes from Defaults ─────────────────────┐\n")
		changes := []string{}

		if c.config.Host != "" {
			changes = append(changes, fmt.Sprintf("│ Host:      %s", c.config.Host))
		}
		if c.config.Port != 0 && c.config.Port != 22 {
			changes = append(changes, fmt.Sprintf("│ Port:      %d (default: 22)", c.config.Port))
		}
		if c.config.BasePath != "" {
			changes = append(changes, fmt.Sprintf("│ Base Path: %s", c.config.BasePath))
		}
		if c.config.RedisAddr != "" && c.config.RedisAddr != "localhost:6379" {
			changes = append(changes, fmt.Sprintf("│ Redis:     %s (default: localhost:6379)", c.config.RedisAddr))
		}
		if c.config.ServerURL != "" {
			changes = append(changes, fmt.Sprintf("│ Server:    %s", c.config.ServerURL))
		}

		if len(changes) == 0 {
			output.WriteString("│ (Using all default settings)\n")
		} else {
			for _, change := range changes {
				output.WriteString(change + "\n")
			}
		}
		output.WriteString("└─────────────────────────────────────────────┘\n\n")

		output.WriteString("Full Configuration:\n")
		output.WriteString(fmt.Sprintf("  Host:      %s\n", c.config.Host))
		output.WriteString(fmt.Sprintf("  Port:      %d\n", c.config.Port))
		output.WriteString(fmt.Sprintf("  Base Path: %s\n", c.config.BasePath))
		output.WriteString(fmt.Sprintf("  Redis:     %s\n", c.config.RedisAddr))
		output.WriteString(fmt.Sprintf("  Server:    %s\n", c.config.ServerURL))
		output.WriteString(fmt.Sprintf("  User:      %s\n\n", c.config.User))

		output.WriteString("Use CLI to modify: ml config set <key> <value>")

		return ConfigLoadedMsg(output.String())
	}
}

// LogsLoadedMsg is sent when logs are loaded
type LogsLoadedMsg string

func (c *Controller) loadLogs(jobName string) tea.Cmd {
	return func() tea.Msg {
		// Placeholder - will stream logs from job
		return LogsLoadedMsg("📜 Logs for " + jobName + "\n\n" +
			"Log streaming will appear here...\n\n" +
			"(Requires API: GET /api/jobs/" + jobName + "/logs?follow=true)")
	}
}

// ExportCompletedMsg is sent when export is complete
type ExportCompletedMsg struct {
	JobName string
	Path    string
}

func (c *Controller) exportJob(jobName string) tea.Cmd {
	return func() tea.Msg {
		// Show export in progress
		return model.StatusMsg{
			Text:  "Exporting " + jobName + "... (anonymized)",
			Level: "info",
		}
	}
}