main.go

package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/charmbracelet/log"
	"github.com/sahilm/fuzzy"
	"gopkg.in/yaml.v3"
)

const version = "1.1.0"

type Shell struct {
	Name    string `yaml:"name"`
	Command string `yaml:"command"`
	Default bool   `yaml:"default"`
}

type Config struct {
	Shells     []Shell `yaml:"shells"`
	ShowSearch bool    `yaml:"show_search"`
}

type model struct {
	shells     []Shell
	filtered   []Shell
	cursor     int
	selected   *Shell
	query      string
	width      int
	height     int
	quitting   bool
	showError  bool
	showSearch bool
}

type clearErrorMsg struct{}

var (
	normalStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("252"))

	selectedStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("30")).
			Bold(true)

	borderStyle = lipgloss.NewStyle().
			Border(lipgloss.RoundedBorder()).
			BorderForeground(lipgloss.Color("240")).
			Padding(1, 3)

	errorBorderStyle = lipgloss.NewStyle().
				Border(lipgloss.RoundedBorder()).
				BorderForeground(lipgloss.Color("124")).
				Padding(1, 3)

	errorTextStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("124")).
			Bold(true)

	dimStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("240"))

	inputStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("252"))
)

func getConfigPath() string {
	configDir := os.Getenv("XDG_CONFIG_HOME")
	if configDir == "" {
		home, err := os.UserHomeDir()
		if err != nil {
			log.Fatal("Could not determine home directory", "error", err)
		}
		configDir = filepath.Join(home, ".config")
	}
	return filepath.Join(configDir, "session-picker", "config.yaml")
}

func ensureConfigDir() error {
	configPath := getConfigPath()
	configDir := filepath.Dir(configPath)
	return os.MkdirAll(configDir, 0755)
}

func loadConfig() (Config, error) {
	var config Config
	configPath := getConfigPath()

	data, err := os.ReadFile(configPath)
	if err != nil {
		return config, err
	}

	err = yaml.Unmarshal(data, &config)
	return config, err
}

func validateConfig(config Config) error {
	if len(config.Shells) == 0 {
		return fmt.Errorf("no shells configured")
	}

	defaultCount := 0
	names := make(map[string]bool)

	for i, shell := range config.Shells {
		if shell.Name == "" {
			return fmt.Errorf("shell at index %d has no name", i)
		}
		if shell.Command == "" {
			return fmt.Errorf("shell '%s' has no command", shell.Name)
		}
		if names[shell.Name] {
			return fmt.Errorf("duplicate shell name: '%s'", shell.Name)
		}
		names[shell.Name] = true
		if shell.Default {
			defaultCount++
		}
	}

	if defaultCount > 1 {
		return fmt.Errorf("multiple shells marked as default (only one allowed)")
	}

	return nil
}

func createDefaultConfig() error {
	if err := ensureConfigDir(); err != nil {
		return err
	}

	configPath := getConfigPath()
	if _, err := os.Stat(configPath); err == nil {
		return nil // Config already exists
	}

	defaultConfig := `# Show the search input box (fuzzy search works regardless)
show_search: false

shells:
  - name: "Local"
    command: "exec $SHELL"
    default: true
  # - name: "Server"
  #   command: "ssh user@server"
`
	return os.WriteFile(configPath, []byte(defaultConfig), 0644)
}

func sortShellsWithDefaultFirst(shells []Shell) []Shell {
	result := make([]Shell, 0, len(shells))
	var others []Shell

	for _, s := range shells {
		if s.Default {
			result = append(result, s)
		} else {
			others = append(others, s)
		}
	}

	return append(result, others...)
}

func initialModel(config Config) model {
	sorted := sortShellsWithDefaultFirst(config.Shells)
	return model{
		shells:     sorted,
		filtered:   sorted,
		cursor:     0,
		showSearch: config.ShowSearch,
	}
}

func (m model) Init() tea.Cmd {
	return nil
}

func clearErrorAfter(d time.Duration) tea.Cmd {
	return tea.Tick(d, func(t time.Time) tea.Msg {
		return clearErrorMsg{}
	})
}

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case clearErrorMsg:
		m.showError = false
		return m, nil

	case tea.WindowSizeMsg:
		m.width = msg.Width
		m.height = msg.Height
		return m, nil

	case tea.KeyMsg:
		switch msg.Type {
		case tea.KeyCtrlC:
			m.quitting = true
			return m, tea.Quit

		case tea.KeyEsc:
			if m.query != "" {
				m.query = ""
				m.filterShells()
				return m, nil
			}
			m.showError = true
			return m, clearErrorAfter(700 * time.Millisecond)

		case tea.KeyEnter:
			if len(m.filtered) > 0 && m.cursor < len(m.filtered) {
				m.selected = &m.filtered[m.cursor]
			}
			return m, tea.Quit

		case tea.KeyUp:
			if m.cursor > 0 {
				m.cursor--
			}
			return m, nil

		case tea.KeyDown:
			if m.cursor < len(m.filtered)-1 {
				m.cursor++
			}
			return m, nil

		case tea.KeyBackspace:
			if len(m.query) > 0 {
				m.query = m.query[:len(m.query)-1]
				m.filterShells()
			}
			return m, nil

		default:
			if msg.Type == tea.KeyRunes {
				m.query += string(msg.Runes)
				m.filterShells()
			}
			return m, nil
		}
	}

	return m, nil
}

func (m *model) filterShells() {
	if m.query == "" {
		m.filtered = m.shells
		m.cursor = 0
		return
	}

	names := make([]string, len(m.shells))
	for i, s := range m.shells {
		names[i] = s.Name
	}

	matches := fuzzy.Find(m.query, names)
	matched := make([]Shell, len(matches))
	for i, match := range matches {
		matched[i] = m.shells[match.Index]
	}

	m.filtered = sortShellsWithDefaultFirst(matched)

	if m.cursor >= len(m.filtered) {
		m.cursor = max(0, len(m.filtered)-1)
	}
}

func (m model) View() string {
	if m.quitting && m.selected == nil {
		return ""
	}

	var b strings.Builder

	// Show search box if enabled
	if m.showSearch && !m.showError {
		if m.query == "" {
			b.WriteString(dimStyle.Render("> type to search..."))
		} else {
			b.WriteString(inputStyle.Render("> " + m.query))
		}
		b.WriteString("\n\n")
	}

	if m.showError {
		b.WriteString(errorTextStyle.Render("ERROR: You have to pick a shell"))
	} else if len(m.filtered) == 0 {
		b.WriteString(dimStyle.Render("no matches"))
	} else {
		for i, shell := range m.filtered {
			style := normalStyle
			if i == m.cursor {
				style = selectedStyle
			}

			b.WriteString(style.Render(shell.Name))
			if i < len(m.filtered)-1 {
				b.WriteString("\n")
			}
		}
	}

	border := borderStyle
	if m.showError {
		border = errorBorderStyle
	}
	content := border.Render(b.String())

	if m.width > 0 && m.height > 0 {
		return lipgloss.Place(
			m.width,
			m.height,
			lipgloss.Center,
			lipgloss.Center,
			content,
		)
	}

	return content
}

func printHelp() {
	help := `sp - Session Picker v` + version + `

A simple, configurable terminal session picker TUI for quickly
connecting to different shells and remote hosts.

Usage:
  sp              Launch the session picker
  sp help         Show this help message
  sp config       Open config file in $EDITOR
  sp config -v    Validate config file
  sp version      Show version

Config location: ` + getConfigPath() + `

Config format:
  show_search: false  # Show/hide the search input box
  shells:
    - name: "Local"
      command: "exec $SHELL"
      default: true
    - name: "Server"
      command: "ssh user@server"

Controls:
  Up/Down    Navigate
  Enter      Select session
  Esc        Clear search / Show error if empty
  Type       Fuzzy search
  Ctrl+C     Quit
`
	fmt.Print(help)
}

func printVersion() {
	fmt.Printf("sp version %s\n", version)
}

func openConfig() {
	editor := os.Getenv("EDITOR")
	if editor == "" {
		editor = "vi"
	}

	configPath := getConfigPath()

	// Create default config if it doesn't exist
	if err := createDefaultConfig(); err != nil {
		log.Error("Failed to create config directory", "error", err)
		os.Exit(1)
	}

	// Create empty config if file doesn't exist
	if _, err := os.Stat(configPath); os.IsNotExist(err) {
		if err := createDefaultConfig(); err != nil {
			log.Error("Failed to create default config", "error", err)
			os.Exit(1)
		}
	}

	cmd := exec.Command(editor, configPath)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		log.Error("Failed to open editor", "error", err)
		os.Exit(1)
	}

	// Validate after editing
	config, err := loadConfig()
	if err != nil {
		log.Error("Failed to load config after editing", "error", err)
		os.Exit(1)
	}

	if err := validateConfig(config); err != nil {
		log.Error("Config validation failed", "error", err)
		os.Exit(1)
	}

	log.Info("Config is valid", "shells", len(config.Shells))
}

func validateConfigCmd() {
	config, err := loadConfig()
	if err != nil {
		log.Error("Failed to load config", "error", err)
		os.Exit(1)
	}

	if err := validateConfig(config); err != nil {
		log.Error("Config validation failed", "error", err)
		os.Exit(1)
	}

	log.Info("Config is valid", "shells", len(config.Shells))
}

func main() {
	args := os.Args[1:]

	if len(args) > 0 {
		switch args[0] {
		case "help", "-h", "--help":
			printHelp()
			return
		case "version", "-v", "--version":
			printVersion()
			return
		case "config":
			if len(args) > 1 && (args[1] == "-v" || args[1] == "--validate") {
				validateConfigCmd()
				return
			}
			openConfig()
			return
		default:
			fmt.Fprintf(os.Stderr, "Unknown command: %s\n", args[0])
			fmt.Fprintln(os.Stderr, "Run 'sp help' for usage.")
			os.Exit(1)
		}
	}

	// Create default config if none exists
	if err := createDefaultConfig(); err != nil {
		log.Error("Failed to create config", "error", err)
		os.Exit(1)
	}

	config, err := loadConfig()
	if err != nil {
		log.Error("Failed to load config", "error", err, "path", getConfigPath())
		log.Info("Run 'sp config' to create a config file")
		os.Exit(1)
	}

	if err := validateConfig(config); err != nil {
		log.Error("Invalid config", "error", err)
		os.Exit(1)
	}

	p := tea.NewProgram(
		initialModel(config),
		tea.WithAltScreen(),
	)

	result, err := p.Run()
	if err != nil {
		log.Error("Failed to run", "error", err)
		os.Exit(1)
	}

	m := result.(model)
	if m.selected != nil {
		shell := os.Getenv("SHELL")
		if shell == "" {
			shell = "/bin/sh"
		}

		args := []string{shell, "-c", m.selected.Command}
		syscall.Exec(shell, args, os.Environ())
	}
}