HostKeeper/pkg/tui/tui.go

package tui

import (
	"context"

	tea "github.com/charmbracelet/bubbletea"

	"git.tukangketik.id/swanadiva/hostkeeper/internal/models"
	"git.tukangketik.id/swanadiva/hostkeeper/pkg/knownhosts"
	"git.tukangketik.id/swanadiva/hostkeeper/pkg/storage"
)

// Screen represents different TUI screens (deprecated, use tabs)
type Screen int

const (
	ScreenHostList Screen = iota
	ScreenConnection
	ScreenSettings
)

// Model represents the main TUI model
type Model struct {
	tabs          *TabManager
	CurrentScreen Screen // deprecated, kept for backward compat
	Hosts         []*models.Host // deprecated
	SelectedIndex int // deprecated
	Error         error
	Quit          bool
	dataDir       string
	program       *tea.Program

	// Security
	storagePassword    string
	knownHosts         *knownhosts.KnownHosts
	showEncryptPrompt  bool
}

// New creates a new TUI model
func New() *Model {
	hostList := NewHostListTab()
	tm := NewTabManager(hostList)

	return &Model{
		tabs:          tm,
		CurrentScreen: ScreenHostList,
		SelectedIndex: 0,
		Quit:          false,
	}
}

// Init initializes the TUI
func (m *Model) Init() tea.Cmd {
	if m.showEncryptPrompt {
		// Show password prompt tab
		tab := NewPasswordPromptTab(passwordModeSetup, m.dataDir, func(password string) {
			m.storagePassword = password
		})
		cmd := m.tabs.Add(tab)
		m.showEncryptPrompt = false
		return cmd
	}
	return m.tabs.Init()
}

// SetProgram stores a reference to the tea.Program for sending messages from goroutines
func (m *Model) SetProgram(p *tea.Program) {
	m.program = p
}

// SetStoragePassword sets the master password for encrypted storage
func (m *Model) SetStoragePassword(password string) {
	m.storagePassword = password
}

// SetKnownHosts sets the known_hosts manager for host key verification
func (m *Model) SetKnownHosts(kh *knownhosts.KnownHosts) {
	m.knownHosts = kh
}

// ShowEncryptPrompt sets a flag to show the encryption setup prompt on first render
func (m *Model) ShowEncryptPrompt() {
	m.showEncryptPrompt = true
}

// Update handles messages and updates the model
func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.KeyMsg:
		// Only hard quit on ctrl+c (let tabs handle q)
		if msg.String() == "ctrl+c" {
			m.Quit = true
			return m, tea.Quit
		}

	case quitMsg:
		m.Quit = true
		return m, tea.Quit

	case sshConnectToMsg:
		return m, tea.Batch(tea.ClearScreen, sshConnectCmd(msg.host, m.dataDir))

	case sshExitMsg:
		if msg.err != nil {
			m.Error = msg.err
		}
		return m, tea.ClearScreen

	case openHostFormMsg:
		var tab Tab
		if msg.editing != nil {
			tab = NewEditHostFormTab(msg.editing, m.dataDir)
		} else {
			tab = NewAddHostFormTab(m.dataDir)
		}
		cmd := m.tabs.Add(tab)
		return m, cmd

	case passwordSetMsg:
		// Password was set — store it and proceed to host list
		m.storagePassword = msg.password
		// Close the password prompt tab, host list should already be there
		return m, nil

	case openEncryptPromptMsg:
		// Show password setup prompt
		tab := NewPasswordPromptTab(passwordModeSetup, m.dataDir, func(password string) {
			m.storagePassword = password
		})
		cmd := m.tabs.Add(tab)
		return m, cmd

	case openSFTPMsg:
		tab := NewSFTPBrowserTab(msg.host, m.dataDir)
		if m.program != nil {
			tab.SetProgram(m.program)
		}
		if m.storagePassword != "" {
			tab.SetStoragePassword(m.storagePassword)
		}
		if m.knownHosts != nil {
			tab.SetPassphraseCallback(func() string {
				// TODO: prompt for passphrase in TUI
				return ""
			})
		}
		cmd := m.tabs.Add(tab)
		return m, cmd

	case openKeyListMsg:
		tab := NewKeyListTab(m.dataDir)
		cmd := m.tabs.Add(tab)
		return m, cmd

	case openSnippetListMsg:
		tab := NewSnippetListTab(m.dataDir)
		cmd := m.tabs.Add(tab)
		return m, cmd

	case openKeyFormMsg:
		var tab Tab
		if msg.editing != nil {
			tab = NewEditKeyFormTab(msg.editing, m.dataDir)
		} else {
			tab = NewAddKeyFormTab(m.dataDir)
		}
		cmd := m.tabs.Add(tab)
		return m, cmd

	case openSnippetFormMsg:
		var tab Tab
		if msg.editing != nil {
			tab = NewEditSnippetFormTab(msg.editing, m.dataDir)
		} else {
			tab = NewAddSnippetFormTab(m.dataDir)
		}
		cmd := m.tabs.Add(tab)
		return m, cmd

	case closeFormMsg:
		if m.tabs.Len() > 1 {
			m.tabs.CloseActive()
		}
		return m, nil

	case saveHostResultMsg:
		if msg.err != nil {
			m.Error = msg.err
			return m, nil
		}
		// Close the form tab, switch back to host list
		if m.tabs.Len() > 1 {
			m.tabs.CloseActive()
		}
		// Reload hosts
		return m, func() tea.Msg {
			store, err := storage.NewJSONStorage(m.dataDir)
			if err != nil {
				return nil
			}
			hosts, err := store.ListHosts(context.Background())
			if err != nil {
				return nil
			}
			return loadedHostsMsg{hosts: hosts}
		}

	case loadedHostsMsg:
		m.LoadHosts(msg.hosts)

	case saveKeyResultMsg:
		if msg.err != nil {
			m.Error = msg.err
			return m, nil
		}
		// Close form tab and switch back to list
		if m.tabs.Len() > 1 {
			m.tabs.CloseActive()
		}
		// Reload keys
		return m, func() tea.Msg {
			store, err := storage.NewJSONStorage(m.dataDir)
			if err != nil {
				return nil
			}
			keys, err := store.ListKeyPairs(context.Background())
			if err != nil {
				return nil
			}
			return keyListLoadedMsg{keys: keys}
		}

	case deleteKeyResultMsg:
		if msg.err != nil {
			m.Error = msg.err
		}
		// Reload keys
		return m, func() tea.Msg {
			store, err := storage.NewJSONStorage(m.dataDir)
			if err != nil {
				return nil
			}
			keys, err := store.ListKeyPairs(context.Background())
			if err != nil {
				return nil
			}
			return keyListLoadedMsg{keys: keys}
		}

	case keyListLoadedMsg:
		if msg.err != nil {
			m.Error = msg.err
			return m, nil
		}
		if kt := FindKeyListTab(m.tabs.tabs); kt != nil {
			kt.SetKeys(msg.keys)
		}

	case saveSnippetResultMsg:
		if msg.err != nil {
			m.Error = msg.err
			return m, nil
		}
		// Close form tab and switch back to list
		if m.tabs.Len() > 1 {
			m.tabs.CloseActive()
		}
		// Reload snippets
		return m, func() tea.Msg {
			store, err := storage.NewJSONStorage(m.dataDir)
			if err != nil {
				return nil
			}
			snippets, err := store.ListSnippets(context.Background())
			if err != nil {
				return nil
			}
			return snippetListLoadedMsg{snippets: snippets}
		}

	case deleteSnippetResultMsg:
		if msg.err != nil {
			m.Error = msg.err
		}
		// Reload snippets
		return m, func() tea.Msg {
			store, err := storage.NewJSONStorage(m.dataDir)
			if err != nil {
				return nil
			}
			snippets, err := store.ListSnippets(context.Background())
			if err != nil {
				return nil
			}
			return snippetListLoadedMsg{snippets: snippets}
		}

	case snippetListLoadedMsg:
		if msg.err != nil {
			m.Error = msg.err
			return m, nil
		}
		if st := FindSnippetListTab(m.tabs.tabs); st != nil {
			st.SetSnippets(msg.snippets)
		}

	}

	cmd, err := m.tabs.Update(msg)
	if err != nil {
		m.Error = err
	}

	// Sync deprecated fields
	if ht := FindHostListTab(m.tabs.tabs); ht != nil {
		m.Hosts = ht.Hosts()
		m.SelectedIndex = ht.SelectedIndex()
	}

	return m, cmd
}

// View renders the TUI
func (m *Model) View() string {
	if m.Quit {
		m.tabs = nil
		return "Thanks for using hostkeeper!\n"
	}

	if m.tabs == nil || m.tabs.Len() == 0 {
		return "No tabs open. Press 'q' to quit.\n"
	}

	// Pass error to host list tab for display
	if m.Error != nil {
		if ht := FindHostListTab(m.tabs.tabs); ht != nil {
			ht.err = m.Error
		}
		m.Error = nil
	}

	return m.tabs.View()
}

// LoadHosts loads hosts into the TUI model
func (m *Model) LoadHosts(hosts []*models.Host) {
	if m.tabs == nil {
		return
	}

	if ht := FindHostListTab(m.tabs.tabs); ht != nil {
		ht.SetHosts(hosts)
		m.Hosts = hosts
	}
}

// SetDataDir sets the data directory for SSH connections
func (m *Model) SetDataDir(dir string) {
	m.dataDir = dir
}

// TabManager returns the underlying tab manager
func (m *Model) TabManager() *TabManager {
	return m.tabs
}