correx/apps/tui-go/internal/app/update.go

package app

import (
	"strings"
	"time"

	tea "github.com/charmbracelet/bubbletea"
	"github.com/correx/tui-go/internal/protocol"
	"github.com/correx/tui-go/internal/ws"
)

// --- Bubble Tea messages bridging the ws channels ---

type serverMsg struct{ m protocol.ServerMessage }
type connMsg struct{ s ws.Status }
type tickMsg struct{}

func readServer(c *ws.Client) tea.Cmd {
	return func() tea.Msg { return serverMsg{<-c.Incoming()} }
}

func readConn(c *ws.Client) tea.Cmd {
	return func() tea.Msg { return connMsg{<-c.Conn()} }
}

func tick() tea.Cmd {
	return tea.Tick(120*time.Millisecond, func(time.Time) tea.Msg { return tickMsg{} })
}

// Init starts the channel readers and the cursor-blink tick.
func (m Model) Init() tea.Cmd {
	return tea.Batch(readServer(m.client), readConn(m.client), tick())
}

// Update is the single reducer. View renders purely from the returned model.
func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		m.width, m.height = msg.Width, msg.Height
		return m, nil

	case tickMsg:
		m.frame++
		return m, tick()

	case connMsg:
		m.applyConn(msg.s)
		return m, readConn(m.client)

	case serverMsg:
		m.applyServerPhased(msg.m)
		return m, readServer(m.client)

	case tea.KeyMsg:
		return m.handleKey(msg)
	}
	return m, nil
}

func (m *Model) applyConn(s ws.Status) {
	switch {
	case s.Connected:
		m.connected = true
		m.reconnecting = false
		m.snapshotPhase = true
		m.pendingEvents = nil
	case s.Reconnecting:
		m.connected = false
		m.reconnecting = true
		m.snapshotPhase = true
	}
}

// applyServerPhased buffers event-bearing messages during the snapshot phase and
// drains them atomically on snapshot_complete.
func (m *Model) applyServerPhased(msg protocol.ServerMessage) {
	if msg.Type == protocol.TypeSnapshotComplete {
		m.snapshotPhase = false
		buffered := m.pendingEvents
		m.pendingEvents = nil
		for _, e := range buffered {
			m.applyServer(e)
		}
		return
	}
	if m.snapshotPhase && msg.IsEventBearing() {
		m.pendingEvents = append(m.pendingEvents, msg)
		return
	}
	m.applyServer(msg)
}

// --- key handling ---

func (m Model) handleKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
	debugLog("KEY type=%v runes=%q alt=%v | ds=%s input=%s edit=%s overlay=%d steering=%v steerBuf=%q dismissed=%v entered=%v sel=%s",
		k.Type, string(k.Runes), k.Alt, m.displayState(), m.inputMode, m.editMode, m.overlay,
		m.steering, m.steerBuffer, m.approvalDismissed, m.sessionEntered, m.selectedID)
	// Ctrl+C is a universal hard-quit safety; everything else is bare-key modal.
	if k.Type == tea.KeyCtrlC {
		m.quitting = true
		return m, tea.Quit
	}
	if m.overlay != OverlayNone {
		return m.handleOverlayKey(k)
	}
	if m.editMode == ModeInsert {
		return m.handleInsertKey(k)
	}
	return m.handleNormalKey(k)
}

// handleNormalKey processes vim-style bare-key commands.
func (m Model) handleNormalKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
	ds := m.displayState()
	switch k.Type {
	case tea.KeyUp:
		m.navUp()
		return m, nil
	case tea.KeyDown:
		m.navDown()
		return m, nil
	case tea.KeyEnter:
		return m.normalEnter()
	case tea.KeyEsc:
		if ds == StateApproval {
			m.approvalDismissed = true
			return m, nil
		}
		m.filter = ""
		return m, nil
	case tea.KeyCtrlA:
		if ds == StateApproval {
			return m.decide("APPROVE")
		}
		return m, nil
	case tea.KeyCtrlR:
		if ds == StateApproval {
			return m.decide("REJECT")
		}
		return m, nil
	case tea.KeyCtrlX:
		if m.currentDiff() != "" {
			m.overlay = OverlayDiff
			m.diffScrollOffset = 0
		}
		return m, nil
	}
	if k.Type != tea.KeyRunes {
		return m, nil
	}
	switch string(k.Runes) {
	case "i":
		m.enterInsert(ModeRouter)
	case "/":
		m.enterInsert(ModeFilter)
	case "j":
		m.navDown()
	case "k":
		m.navUp()
	case "p":
		m.overlay = OverlayPalette
		m.paletteFilter = ""
		m.paletteIndex = 0
	case "e":
		m.overlay = OverlayEventInspector
		m.overlayEventIdx = 0
	case "t":
		m.overlay = OverlayToolPalette
	case "v":
		m.openArtifacts()
	case "g":
		m.openConfig()
	case "m":
		m.openModelsOverlay()
	case "w":
		if ds == StateIdle {
			m.wfVisible = !m.wfVisible
			if m.wfVisible {
				m.wfIndex = 0
			} else {
				m.wfIndex = -1
			}
		}
	case "l":
		if ds != StateIdle {
			m.sessionEntered = false
			m.approvalDismissed = false
		}
	case "s":
		if ds == StateApproval {
			m.steering = true
			m.editMode = ModeInsert
		} else {
			m.cycleChatMode()
		}
	case "c":
		if m.selectedID != "" {
			m.client.Send(protocol.CancelSession(m.selectedID))
		}
	case "a":
		if ds == StateApproval {
			return m.decide("APPROVE")
		}
		if ds == StateInSession {
			if s := m.session(m.selectedID); s != nil && s.Pending != nil {
				m.approvalDismissed = false
			}
		}
	case "r":
		if ds == StateApproval {
			return m.decide("REJECT")
		}
	case "A":
		if ds == StateApproval {
			return m.autoApprove()
		}
	case "q":
		m.quitting = true
		return m, tea.Quit
	}
	return m, nil
}

// handleInsertKey processes typing (chat/filter/steer note).
func (m Model) handleInsertKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
	if m.steering {
		switch k.Type {
		case tea.KeyEsc:
			m.steering = false
			m.editMode = ModeNormal
		case tea.KeyEnter:
			return m.submitApproval()
		case tea.KeyBackspace:
			if n := len(m.steerBuffer); n > 0 {
				m.steerBuffer = m.steerBuffer[:n-1]
			}
		case tea.KeyRunes, tea.KeySpace:
			m.steerBuffer += string(k.Runes)
		}
		return m, nil
	}

	switch k.Type {
	case tea.KeyEsc:
		m.editMode = ModeNormal
		if m.inputMode == ModeFilter {
			m.inputMode = ModeRouter
		}
		if m.inputMode == ModeIntent {
			m.wfPendingID = ""
			m.wfPendingName = ""
			m.clearInput()
			m.inputMode = ModeRouter
		}
	case tea.KeyEnter:
		return m.submit()
	case tea.KeyBackspace:
		m.backspace()
		m.syncFilter()
	case tea.KeyLeft:
		if m.inputCursor > 0 {
			m.inputCursor--
		}
	case tea.KeyRight:
		if m.inputCursor < len(m.inputBuffer) {
			m.inputCursor++
		}
	case tea.KeyUp:
		if m.inputMode == ModeFilter {
			m.listNav(-1)
		} else {
			m.historyPrev()
		}
	case tea.KeyDown:
		if m.inputMode == ModeFilter {
			m.listNav(1)
		} else {
			m.historyNext()
		}
	case tea.KeyRunes, tea.KeySpace:
		m.appendRunes(string(k.Runes))
		m.syncFilter()
	}
	return m, nil
}

// beginIntent stashes the chosen workflow and switches to a text-entry prompt for the
// freeform request. Submitting sends StartSession(id, intent); an empty line starts with no
// intent (fixed-task workflows). Esc cancels.
func (m *Model) beginIntent(wf Workflow) {
	m.wfPendingID = wf.ID
	m.wfPendingName = wf.ID
	m.wfVisible = false
	m.wfIndex = -1
	m.clearInput()
	m.editMode = ModeInsert
	m.inputMode = ModeIntent
}

func (m *Model) enterInsert(mode InputMode) {
	m.editMode = ModeInsert
	m.inputMode = mode
	if mode == ModeFilter {
		m.inputBuffer = m.filter
		m.inputCursor = len(m.inputBuffer)
	}
}

// syncFilter keeps the live session filter in step with the input buffer.
func (m *Model) syncFilter() {
	if m.inputMode == ModeFilter {
		m.filter = m.inputBuffer
	}
}

func (m Model) normalEnter() (tea.Model, tea.Cmd) {
	if m.displayState() != StateIdle {
		return m, nil
	}
	if m.wfVisible && m.wfIndex >= 0 && m.wfIndex < len(m.workflows) {
		m.beginIntent(m.workflows[m.wfIndex])
		return m, nil
	}
	if m.selectedID != "" {
		m.sessionEntered = true
	}
	return m, nil
}

func (m Model) autoApprove() (tea.Model, tea.Cmd) {
	s := m.session(m.selectedID)
	if s == nil || s.Pending == nil {
		return m, nil
	}
	p := s.Pending
	m.client.Send(protocol.CreateGrant(p.SessionID, "SESSION", []string{p.Tier}, "auto-approved via TUI", p.ToolName))
	m.client.Send(protocol.ApprovalResponse(p.RequestID, "APPROVE", nil))
	s.Pending = nil
	m.steerBuffer = ""
	m.steering = false
	m.approvalDismissed = false
	return m, nil
}

func (m Model) handleOverlayKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
	// Config overlay owns all its keys (esc cancels an in-progress edit rather than closing).
	if m.overlay == OverlayConfig {
		return m.handleConfigKey(k)
	}
	if k.Type == tea.KeyEsc {
		m.overlay = OverlayNone
		return m, nil
	}
	switch m.overlay {
	case OverlayPalette:
		return m.handlePaletteKey(k)
	case OverlayDiff:
		switch {
		case k.Type == tea.KeyUp || runeIs(k, "k"):
			if m.diffScrollOffset > 0 {
				m.diffScrollOffset--
			}
		case k.Type == tea.KeyDown || runeIs(k, "j"):
			if m.diffScrollOffset < m.diffMaxScroll() {
				m.diffScrollOffset++
			}
		case k.Type == tea.KeyCtrlX:
			m.overlay = OverlayNone
		}
	case OverlayEventInspector:
		evs := m.currentEvents()
		switch {
		case k.Type == tea.KeyUp || runeIs(k, "k"):
			if m.overlayEventIdx > 0 {
				m.overlayEventIdx--
			}
		case k.Type == tea.KeyDown || runeIs(k, "j"):
			if m.overlayEventIdx < len(evs)-1 {
				m.overlayEventIdx++
			}
		}
	case OverlayToolPalette:
		if runeIs(k, "t") {
			m.overlay = OverlayNone
		}
	case OverlayArtifacts:
		switch {
		case k.Type == tea.KeyUp || runeIs(k, "k"):
			if m.artifactsIndex > 0 {
				m.artifactsIndex--
				m.artifactScroll = 0
			}
		case k.Type == tea.KeyDown || runeIs(k, "j"):
			if m.artifactsIndex < len(m.artifacts)-1 {
				m.artifactsIndex++
				m.artifactScroll = 0
			}
		case k.Type == tea.KeyPgUp:
			if m.artifactScroll > 0 {
				m.artifactScroll--
			}
		case k.Type == tea.KeyPgDown:
			contentW := m.modalWidth() - 6
			if m.artifactScroll < m.artifactContentMaxScroll(contentW) {
				m.artifactScroll++
			}
		case runeIs(k, "v"):
			m.overlay = OverlayNone
		}
	case OverlayModels:
		switch {
		case k.Type == tea.KeyUp || runeIs(k, "k"):
			if m.modelsIndex > 0 {
				m.modelsIndex--
			}
		case k.Type == tea.KeyDown || runeIs(k, "j"):
			if m.modelsIndex < len(m.availableModels)-1 {
				m.modelsIndex++
			}
		case k.Type == tea.KeyEnter:
			if m.modelsIndex >= 0 && m.modelsIndex < len(m.availableModels) {
				m.client.Send(protocol.SwapModel(m.availableModels[m.modelsIndex]))
				m.overlay = OverlayNone
			}
		case runeIs(k, "c"):
			m.client.Send(protocol.ClearModelPin())
			m.overlay = OverlayNone
		case runeIs(k, "m"):
			m.overlay = OverlayNone
		}
	}
	return m, nil
}

// openArtifacts opens the artifact viewer for the selected session and requests its
// artifacts from the server. No-op when no session is selected.
func (m *Model) openArtifacts() {
	if m.selectedID == "" {
		return
	}
	m.overlay = OverlayArtifacts
	m.artifactsIndex = 0
	m.artifactScroll = 0
	// Reuse a cached listing only if it's for this session; otherwise show a loading state.
	if m.artifactsFor != m.selectedID {
		m.artifacts = nil
		m.artifactsLoading = true
	}
	m.client.Send(protocol.ListArtifacts(m.selectedID))
}

// openModelsOverlay opens the model picker, pre-selecting the resident model.
func (m *Model) openModelsOverlay() {
	m.overlay = OverlayModels
	m.modelsIndex = 0
	for i, id := range m.availableModels {
		if id == m.currentModel {
			m.modelsIndex = i
			break
		}
	}
}

func runeIs(k tea.KeyMsg, s string) bool {
	return k.Type == tea.KeyRunes && string(k.Runes) == s
}

func (m Model) handlePaletteKey(k tea.KeyMsg) (tea.Model, tea.Cmd) {
	cmds := m.filteredPalette()
	switch k.Type {
	case tea.KeyUp:
		if m.paletteIndex > 0 {
			m.paletteIndex--
		}
	case tea.KeyDown:
		if m.paletteIndex < len(cmds)-1 {
			m.paletteIndex++
		}
	case tea.KeyEnter:
		if m.paletteIndex >= 0 && m.paletteIndex < len(cmds) {
			return m.execPalette(cmds[m.paletteIndex].id)
		}
	case tea.KeyBackspace:
		if n := len(m.paletteFilter); n > 0 {
			m.paletteFilter = m.paletteFilter[:n-1]
			m.paletteIndex = 0
		}
	case tea.KeyRunes, tea.KeySpace:
		m.paletteFilter += string(k.Runes)
		m.paletteIndex = 0
	}
	return m, nil
}

type paletteCmd struct{ id, title, hint string }

func paletteCommands() []paletteCmd {
	return []paletteCmd{
		{"workflows", "start workflow", "open the workflow picker"},
		{"tools", "tool palette", "tools for the current stage"},
		{"models", "swap model", "pick / pin the local model"},
		{"events", "event inspector", "browse the event stream"},
		{"artifacts", "view artifacts", "browse this session's artifacts"},
		{"config", "edit config", "view / change correx settings"},
		{"mode", "toggle mode", "switch chat / steering"},
		{"cancel", "cancel session", "stop the selected session"},
		{"back", "back to list", "leave the current session"},
		{"quit", "quit", "exit correx"},
	}
}

func (m Model) filteredPalette() []paletteCmd {
	f := strings.ToLower(m.paletteFilter)
	if f == "" {
		return paletteCommands()
	}
	var out []paletteCmd
	for _, c := range paletteCommands() {
		if strings.Contains(strings.ToLower(c.title+" "+c.hint), f) {
			out = append(out, c)
		}
	}
	return out
}

func (m Model) execPalette(id string) (tea.Model, tea.Cmd) {
	m.overlay = OverlayNone
	switch id {
	case "workflows":
		m.wfVisible = true
		m.wfIndex = 0
	case "tools":
		m.overlay = OverlayToolPalette
	case "models":
		m.openModelsOverlay()
	case "events":
		m.overlay = OverlayEventInspector
		m.overlayEventIdx = 0
	case "artifacts":
		m.openArtifacts()
	case "config":
		m.openConfig()
	case "mode":
		m.cycleChatMode()
	case "cancel":
		if m.selectedID != "" {
			m.client.Send(protocol.CancelSession(m.selectedID))
		}
	case "back":
		m.sessionEntered = false
		m.approvalDismissed = false
	case "quit":
		m.quitting = true
		return m, tea.Quit
	}
	return m, nil
}

// --- input editing ---

func (m *Model) appendRunes(s string) {
	b := m.inputBuffer
	c := m.inputCursor
	if c > len(b) {
		c = len(b)
	}
	m.inputBuffer = b[:c] + s + b[c:]
	m.inputCursor = c + len(s)
}

func (m *Model) backspace() {
	if m.inputCursor == 0 || len(m.inputBuffer) == 0 {
		return
	}
	c := m.inputCursor
	m.inputBuffer = m.inputBuffer[:c-1] + m.inputBuffer[c:]
	m.inputCursor = c - 1
}

func (m *Model) cycleChatMode() {
	if m.chatMode == ChatModeChat {
		m.chatMode = ChatModeSteering
	} else {
		m.chatMode = ChatModeChat
	}
}

func (m *Model) clearInput() {
	m.inputBuffer = ""
	m.inputCursor = 0
	m.historyIndex = -1
}

// --- submit (IDLE -> StartChat / StartSession, IN_SESSION -> ChatInput) ---

func (m Model) submit() (tea.Model, tea.Cmd) {
	ds := m.displayState()
	if m.inputMode == ModeFilter {
		m.filter = m.inputBuffer
		m.clearInput()
		m.inputMode = ModeRouter
		m.editMode = ModeNormal
		return m, nil
	}
	text := strings.TrimSpace(m.inputBuffer)

	switch ds {
	case StateIdle:
		// Intent line for a workflow being started.
		if m.wfPendingID != "" {
			id := m.wfPendingID
			m.client.Send(protocol.StartSession(id, text))
			m.wfPendingID = ""
			m.wfPendingName = ""
			m.clearInput()
			m.inputMode = ModeRouter
			m.editMode = ModeNormal
			m.pendingWorkflowFocus = true
			return m, nil
		}
		// Workflow picker selection → prompt for the request first.
		if m.wfIndex >= 0 && m.wfIndex < len(m.workflows) {
			m.beginIntent(m.workflows[m.wfIndex])
			return m, nil
		}
		// Entering an already-selected session (blank submit).
		if text == "" && m.selectedID != "" {
			m.sessionEntered = true
			m.clearInput()
			return m, nil
		}
		if text == "" {
			return m, nil
		}
		// New chat session. The client owns the id, so create the entry locally and
		// focus it; the USER/ROUTER turns arrive as chat.turn events (no optimistic echo).
		id := newSessionID()
		s := m.ensureSession(id)
		s.WorkflowID = "chat"
		s.Name = "chat"
		m.selectedID = id
		m.sessionEntered = true
		m.client.Send(protocol.StartChatSession(id, text))
		m.clearInput()
		return m, nil

	case StateInSession:
		if text == "" || m.selectedID == "" {
			return m, nil
		}
		sid := m.selectedID
		hist := m.history[sid]
		hist = append(hist, text)
		if len(hist) > 50 {
			hist = hist[len(hist)-50:]
		}
		m.history[sid] = hist
		// No optimistic echo — the USER turn arrives as a chat.turn event.
		m.client.Send(protocol.ChatInput(sid, text, m.chatMode))
		m.clearInput()
		return m, nil
	}
	return m, nil
}

func (m Model) decide(decision string) (tea.Model, tea.Cmd) {
	s := m.session(m.selectedID)
	if s == nil || s.Pending == nil {
		return m, nil
	}
	var note *string
	if strings.TrimSpace(m.steerBuffer) != "" {
		n := m.steerBuffer
		note = &n
	}
	m.client.Send(protocol.ApprovalResponse(s.Pending.RequestID, decision, note))
	s.Pending = nil
	m.steerBuffer = ""
	m.steering = false
	m.editMode = ModeNormal
	m.approvalDismissed = false
	return m, nil
}

func (m Model) submitApproval() (tea.Model, tea.Cmd) {
	// Enter approves by default; steering note (if any) rides along.
	return m.decide("APPROVE")
}

// --- navigation (list / history) ---

func (m *Model) navUp() {
	if m.displayState() == StateInSession && m.inputMode == ModeRouter {
		m.historyPrev()
		return
	}
	if m.wfVisible {
		m.wfNav(-1)
		return
	}
	if m.displayState() == StateIdle {
		m.listNav(-1)
	}
}

func (m *Model) navDown() {
	if m.displayState() == StateInSession && m.inputMode == ModeRouter {
		m.historyNext()
		return
	}
	if m.wfVisible {
		m.wfNav(1)
		return
	}
	if m.displayState() == StateIdle {
		m.listNav(1)
	}
}

func (m *Model) wfNav(dir int) {
	n := len(m.workflows)
	if n == 0 {
		return
	}
	m.wfIndex = (m.wfIndex + dir + n) % n
}

func (m *Model) listNav(dir int) {
	list := m.filteredSessions()
	if len(list) == 0 {
		return
	}
	idx := -1
	for i, s := range list {
		if s.ID == m.selectedID {
			idx = i
			break
		}
	}
	// No current selection (e.g. nothing entered yet): start at the first
	// session rather than wrapping off the end into an arbitrary entry.
	if idx < 0 {
		m.selectedID = list[0].ID
		m.wfIndex = -1
		return
	}
	n := idx + dir
	if n < 0 {
		n = len(list) - 1
	}
	if n >= len(list) {
		n = 0
	}
	if m.selectedID != list[n].ID {
		m.bgUpdates = 0
	}
	m.selectedID = list[n].ID
	m.wfIndex = -1
}

func (m *Model) historyPrev() {
	hist := m.history[m.selectedID]
	if len(hist) == 0 {
		return
	}
	switch m.historyIndex {
	case -1:
		m.savedBuffer = m.inputBuffer
		m.historyIndex = len(hist) - 1
	case 0:
		return
	default:
		m.historyIndex--
	}
	m.inputBuffer = hist[m.historyIndex]
	m.inputCursor = len(m.inputBuffer)
}

func (m *Model) historyNext() {
	hist := m.history[m.selectedID]
	switch {
	case m.historyIndex == -1:
		return
	case m.historyIndex == len(hist)-1:
		m.historyIndex = -1
		m.inputBuffer = m.savedBuffer
	default:
		m.historyIndex++
		m.inputBuffer = hist[m.historyIndex]
	}
	m.inputCursor = len(m.inputBuffer)
}

func (m *Model) appendRouter(sid string, e RouterEntry) {
	m.routerMessages[sid] = append(m.routerMessages[sid], e)
}

// currentDiff returns the most recent tool diff in the selected transcript.
func (m Model) currentDiff() string {
	if s := m.session(m.selectedID); s != nil && s.Pending != nil && s.Pending.Preview != "" {
		return s.Pending.Preview
	}
	msgs := m.routerMessages[m.selectedID]
	for i := len(msgs) - 1; i >= 0; i-- {
		if msgs[i].Role == "tool" {
			return msgs[i].Content
		}
	}
	return ""
}

func (m Model) currentEvents() []EventEntry {
	if s := m.session(m.selectedID); s != nil {
		return s.Events
	}
	return nil
}