lazyollama.go

package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"os/exec"
	"strings"
	"time"

	"github.com/jesseduffield/gocui"
)

const (
	MODELS_VIEW = "models"
	FOOTER_VIEW = "footer"
)

type Model struct {
	Name     string
	Size     string
	Modified string
	Running  bool
}

type App struct {
	gui           *gocui.Gui
	models        []Model
	selectedModel int
	ctx           context.Context
	cancel        context.CancelFunc
}

func main() {
	app := NewApp()
	if err := app.Run(); err != nil {
		log.Printf("Error running lazyollama: %v", err)
	}
}

func NewApp() *App {
	ctx, cancel := context.WithCancel(context.Background())
	return &App{
		ctx:    ctx,
		cancel: cancel,
	}
}

func (a *App) Run() error {
	g, err := gocui.NewGui(gocui.NewGuiOpts{
		OutputMode: gocui.OutputNormal,
	})
	if err != nil {
		return fmt.Errorf("failed to create GUI: %w", err)
	}
	defer g.Close()

	a.gui = g
	g.Cursor = false
	g.Mouse = true

	g.SetManagerFunc(a.layout)

	// Load initial data
	a.refreshModels()

	// Start background refresh
	go a.backgroundRefresh()

	if err := g.MainLoop(); err != nil && !gocui.IsQuit(err) {
		return err
	}

	a.cancel()
	return nil
}

func (a *App) layout(g *gocui.Gui) error {
	maxX, maxY := g.Size()

	// Models panel (full screen minus footer)
	if v, err := g.SetView(MODELS_VIEW, 0, 0, maxX-1, maxY-2, 0); err != nil {
		if !gocui.IsUnknownView(err) {
			return err
		}
		v.Title = " Models "
		v.Highlight = true
		v.SelBgColor = gocui.ColorBlue
		v.SelFgColor = gocui.ColorWhite
		// Set rounded border using custom frame runes
		v.FrameRunes = []rune{'─', '│', '╭', '╮', '╰', '╯'}
		a.updateModelsView()
		
		if _, err := g.SetCurrentView(MODELS_VIEW); err != nil {
			return err
		}
		
		// Set keybindings
		a.setKeybindings()
	}

	// Footer (borderless instructions)  
	if v, err := g.SetView(FOOTER_VIEW, 0, maxY-2, maxX-1, maxY, 0); err != nil {
		if !gocui.IsUnknownView(err) {
			return err
		}
		v.Frame = false // Remove border
		v.BgColor = gocui.ColorDefault
		v.FgColor = gocui.ColorWhite
	}
	a.updateFooter()

	return nil
}

func (a *App) setKeybindings() error {
	g := a.gui

	// Global keybindings
	g.SetKeybinding("", gocui.KeyCtrlC, gocui.ModNone, a.quit)
	g.SetKeybinding("", 'q', gocui.ModNone, a.quit)
	g.SetKeybinding("", 'r', gocui.ModNone, a.refresh)

	// Models view keybindings
	g.SetKeybinding(MODELS_VIEW, gocui.KeyArrowDown, gocui.ModNone, a.nextModel)
	g.SetKeybinding(MODELS_VIEW, 'j', gocui.ModNone, a.nextModel)
	g.SetKeybinding(MODELS_VIEW, gocui.KeyArrowUp, gocui.ModNone, a.prevModel)
	g.SetKeybinding(MODELS_VIEW, 'k', gocui.ModNone, a.prevModel)
	g.SetKeybinding(MODELS_VIEW, gocui.KeyEnter, gocui.ModNone, a.chatWithModel)
	g.SetKeybinding(MODELS_VIEW, 'x', gocui.ModNone, a.stopModel)

	return nil
}

func (a *App) refreshModels() {
	models, err := a.getOllamaModels()
	if err != nil {
		return
	}
	a.models = models
	a.updateModelsView()
}

func (a *App) getOllamaModels() ([]Model, error) {
	// Get list of models
	cmd := exec.Command("ollama", "list")
	output, err := cmd.Output()
	if err != nil {
		return nil, err
	}

	lines := strings.Split(string(output), "\n")
	var models []Model

	// Skip header line
	for i, line := range lines {
		if i == 0 || strings.TrimSpace(line) == "" {
			continue
		}
		
		fields := strings.Fields(line)
		if len(fields) >= 3 {
			model := Model{
				Name:     fields[0],
				Size:     fields[1],
				Modified: strings.Join(fields[2:], " "),
			}
			models = append(models, model)
		}
	}

	// Check running models
	runningModels := a.getRunningModels()
	for i, model := range models {
		for _, running := range runningModels {
			if model.Name == running {
				models[i].Running = true
				break
			}
		}
	}

	return models, nil
}

func (a *App) getRunningModels() []string {
	cmd := exec.Command("ollama", "ps")
	output, err := cmd.Output()
	if err != nil {
		return nil
	}

	lines := strings.Split(string(output), "\n")
	var running []string

	for i, line := range lines {
		if i == 0 || strings.TrimSpace(line) == "" {
			continue
		}
		
		fields := strings.Fields(line)
		if len(fields) > 0 {
			running = append(running, fields[0])
		}
	}

	return running
}

func (a *App) updateModelsView() {
	v, err := a.gui.View(MODELS_VIEW)
	if err != nil {
		return
	}

	v.Clear()
	for _, model := range a.models {
		status := "○"
		if model.Running {
			status = "●"
		}
		fmt.Fprintf(v, " %s %s\n", status, model.Name)
	}
	
	// Set cursor to selected model
	if len(a.models) > 0 && a.selectedModel < len(a.models) {
		v.SetCursor(0, a.selectedModel)
	}
}

func (a *App) updateFooter() {
	v, err := a.gui.View(FOOTER_VIEW)
	if err != nil {
		return
	}

	v.Clear()
	if len(a.models) > 0 && a.selectedModel < len(a.models) && a.models[a.selectedModel].Running {
		fmt.Fprintf(v, " [j/k] navigate [Enter] chat [x] stop [r] refresh [q] quit")
	} else {
		fmt.Fprintf(v, " [j/k] navigate [Enter] chat [r] refresh [q] quit")
	}
}



func (a *App) backgroundRefresh() {
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-a.ctx.Done():
			return
		case <-ticker.C:
			a.gui.Update(func(g *gocui.Gui) error {
				a.refreshModels()
				return nil
			})
		}
	}
}

// Event handlers
func (a *App) quit(g *gocui.Gui, v *gocui.View) error {
	return gocui.ErrQuit
}

func (a *App) refresh(g *gocui.Gui, v *gocui.View) error {
	a.refreshModels()
	return nil
}

func (a *App) nextModel(g *gocui.Gui, v *gocui.View) error {
	if len(a.models) > 0 && a.selectedModel < len(a.models)-1 {
		a.selectedModel++
		a.updateModelsView()
		a.updateFooter()
	}
	return nil
}

func (a *App) prevModel(g *gocui.Gui, v *gocui.View) error {
	if len(a.models) > 0 && a.selectedModel > 0 {
		a.selectedModel--
		a.updateModelsView()
		a.updateFooter()
	}
	return nil
}

func (a *App) chatWithModel(g *gocui.Gui, v *gocui.View) error {
	if len(a.models) == 0 || a.selectedModel >= len(a.models) {
		return nil
	}

	model := a.models[a.selectedModel]
	
	// Temporarily close the GUI for chat
	g.Close()
	
	// Run ollama with proper stdin/stdout/stderr
	cmd := exec.Command("ollama", "run", model.Name)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout  
	cmd.Stderr = os.Stderr
	cmd.Run()
	
	// Restart the GUI after chat
	newGui, err := gocui.NewGui(gocui.NewGuiOpts{
		OutputMode: gocui.OutputNormal,
	})
	if err != nil {
		return err
	}
	
	a.gui = newGui
	a.gui.Cursor = false
	a.gui.Mouse = true
	a.gui.SetManagerFunc(a.layout)
	
	// Refresh models and restart
	a.refreshModels()
	
	return a.gui.MainLoop()
}

func (a *App) stopModel(g *gocui.Gui, v *gocui.View) error {
	if len(a.models) == 0 || a.selectedModel >= len(a.models) {
		return nil
	}

	model := a.models[a.selectedModel]
	if !model.Running {
		return nil
	}
	
	go func() {
		cmd := exec.Command("ollama", "stop", model.Name)
		if err := cmd.Run(); err == nil {
			a.gui.Update(func(g *gocui.Gui) error {
				a.refreshModels()
				return nil
			})
		}
	}()

	return nil
}