gust / model.go

package gust

import (
	"bufio"
	"context"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/fsnotify/fsnotify"
)

// CommandProcess represents a running command process
type CommandProcess struct {
	cmd          *exec.Cmd
	ctx          context.Context
	cancel       context.CancelFunc
	stdoutReader io.ReadCloser
	stderrReader io.ReadCloser
	stdoutChan   chan stdoutMsg
	stderrChan   chan stderrMsg
	done         chan struct{}
	waitDone     sync.Once
}

// NewCommandProcess creates a new command process
func NewCommandProcess(cmdName string, args ...string) (*CommandProcess, error) {
	ctx, cancel := context.WithCancel(context.Background())
	cmd := exec.CommandContext(ctx, cmdName, args...)

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		cancel()
		return nil, err
	}

	stderrPipe, err := cmd.StderrPipe()
	if err != nil {
		cancel()
		return nil, err
	}

	process := &CommandProcess{
		cmd:          cmd,
		ctx:          ctx,
		cancel:       cancel,
		stdoutReader: stdoutPipe,
		stderrReader: stderrPipe,
		stdoutChan:   make(chan stdoutMsg, 1024),
		stderrChan:   make(chan stderrMsg, 1024),
		done:         make(chan struct{}),
	}

	return process, nil
}

// Start starts the command and output scanners
func (p *CommandProcess) Start() error {
	if err := p.cmd.Start(); err != nil {
		p.cancel()
		return err
	}

	// Start stdout scanner
	go func() {
		seq := 0
		scanner := bufio.NewScanner(p.stdoutReader)
		for scanner.Scan() {
			select {
			case <-p.ctx.Done():
				return
			case p.stdoutChan <- stdoutMsg{line: scanner.Text() + "\n", seq: seq}:
				seq++
				// Line sent
			}
		}
	}()

	// Start stderr scanner
	go func() {
		seq := 0
		scanner := bufio.NewScanner(p.stderrReader)
		for scanner.Scan() {
			select {
			case <-p.ctx.Done():
				return
			case p.stderrChan <- stderrMsg{line: scanner.Text() + "\n", seq: seq}:
				seq++
				// Line sent
			}
		}
	}()

	// Wait for process in background
	go func() {
		p.waitDone.Do(func() {
			p.cmd.Wait()
			close(p.done)
		})
	}()

	return nil
}

func (p *CommandProcess) Stop() {
	if p.cmd == nil || p.cmd.Process == nil {
		return
	}

	p.cancel()
}

// CheckStatus returns the process state if available
func (p *CommandProcess) CheckStatus() *os.ProcessState {
	if p == nil || p.cmd == nil {
		return nil
	}
	return p.cmd.ProcessState
}

// Model represents the application state
type Model struct {
	stdio []IoLine

	messages []CompilerMessage
	extra    []string // whatever was not parsed from the compiler output

	fs     Fs
	w      *fsnotify.Watcher
	config Config

	status       Status
	systemErrors []error

	process *CommandProcess
	start   *time.Time
	end     *time.Time

	viewport      viewport.Model
	ready         bool
	width, height int
}

type Io []IoLine
type IoLine struct {
	Kind IoKind
	Line string
	Seq  int
}
type IoKind int

const (
	Stdout IoKind = iota
	Stderr
)

func (m Model) stdfilter(filter func(kind IoKind) bool) string {
	var lines []stderrMsg

	for _, l := range m.stdio {
		if filter(l.Kind) {
			lines = append(lines, stderrMsg{
				line: l.Line,
				seq:  l.Seq,
			})
		}
	}

	sort.Slice(lines, func(i, j int) bool {
		return lines[i].seq < lines[j].seq
	})

	var b strings.Builder

	for _, l := range lines {
		b.WriteString(l.line)
	}

	return b.String()
}

func (m Model) stderr() string {
	return m.stdfilter(func(kind IoKind) bool {
		return kind == Stderr
	})
}

func (m Model) stdout() string {
	return m.stdfilter(func(kind IoKind) bool {
		return kind == Stdout
	})
}

func NewModel(watcher *fsnotify.Watcher, config Config, options ...func(*Config)) Model {
	for _, o := range options {
		o(&config)
	}

	return Model{
		w:      watcher,
		status: NewStatus(),
		config: config,
	}
}

func (m Model) Init() tea.Cmd {
	return tea.Batch(reload, m.status.Init())
}

func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	var cmd tea.Cmd

	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		headerHeight := lipgloss.Height(m.viewStatus())
		footerHeight := lipgloss.Height(m.viewHelp())
		verticalMarginHeight := headerHeight + footerHeight

		m.width, m.height = msg.Width, msg.Height-verticalMarginHeight
		if !m.ready {
			m.viewport = viewport.New(m.width, m.height)
			m.viewport.YPosition = headerHeight
			m.ready = true
		} else {
			m.viewport.Width, m.viewport.Height = m.width, m.height
		}

		return m, updateViewport
	case tea.KeyMsg:
		switch msg.String() {
		case "q", "ctrl+c":
			if m.process != nil {
				m.process.Stop()
			}
			return m, tea.Quit
		case "r":
			var cmd tea.Cmd
			if m.config.Mode != "run" {
				m.config.Mode = "run"
				cmd = reload
			}
			return m, cmd
		case "b":
			var cmd tea.Cmd
			if m.config.Mode != "build" {
				m.config.Mode = "build"
				cmd = reload
			}
			return m, cmd
		case "s":
			m.config.Summarized = !m.config.Summarized
			return m, nil
		case "o":
			if m.config.Stream == "stdout" {
				m.config.Stream = "stderr"
			} else {
				m.config.Stream = "stdout"
			}
			return m, nil
		case "h", "?":
			m.config.Help = !m.config.Help
			return m, nil
		default:
			m.viewport, cmd = m.viewport.Update(msg)
			return m, cmd
		}
	case StatusKind, spinner.TickMsg:
		var cmd tea.Cmd
		m.status, cmd = m.status.Update(msg)
		return m, tea.Batch(cmd, m.listen())
	case modifiedMsg, reloadMsg:
		m, cmd := m.reload()
		return m, cmd
	case stdoutMsg:
		m.stdio = append(m.stdio, IoLine{
			Kind: Stdout,
			Line: msg.line,
			Seq:  msg.seq,
		})
		return m, m.checkStdout()
	case stderrMsg:
		m.stdio = append(m.stdio, IoLine{
			Kind: Stderr,
			Line: msg.line,
			Seq:  msg.seq,
		})
		messages, extra := Parse(m.stderr())
		fs := Fs{}
		if !m.config.Summarized {
			fs = BuildFs(messages)
			fs.PopulateContext(messages, m.config.Context)
		}
		m.messages = messages
		m.extra = extra
		m.fs = fs
		return m, m.checkStderr()
	case processState:
		end := time.Now()
		if msg.state == nil {
			return m, tea.Batch(m.checkCmdStatus(), m.checkStderr(), m.checkStdout(), updateViewport)
		} else if msg.state.Success() {
			m.end = &end
			return m, Success.Cmd()
		} else if msg.state.Exited() {
			m.end = &end
			return m, Error.Cmd()
		} else {
			return m, tea.Batch(m.checkCmdStatus(), m.checkStderr(), m.checkStdout(), updateViewport)
		}
	case updateViewportMsg:
		m.viewport.SetContent(m.viewBody())
		if m.config.Mode == "build" {
			m.viewport.GotoTop()
		} else {
			m.viewport.GotoBottom()
		}
		return m, nil
	default:
		return m, nil
	}
}

type reloadMsg struct{}

func reload() tea.Msg {
	return reloadMsg{}
}

type updateViewportMsg struct{}

func updateViewport() tea.Msg {
	return updateViewportMsg{}
}

type modifiedMsg struct{}
type watcherrMsg struct{}

func (m Model) listen() tea.Cmd {
	return func() tea.Msg {
		select {
		case event, ok := <-m.w.Events:
			if !ok {
				return watcherrMsg{}
			}

			if event.Has(fsnotify.Write|fsnotify.Create|fsnotify.Remove) &&
				filepath.Ext(event.Name) == ".go" {
				return modifiedMsg{}
			}

			return nil

		case _, ok := <-m.w.Errors:
			if !ok {
				return watcherrMsg{}
			}
			return watcherrMsg{}
		default:
			return nil
		}
	}
}

func (m Model) reload() (Model, tea.Cmd) {
	status := Running.Cmd()
	m.systemErrors = []error{}

	// Clean up old process
	if m.process != nil {
		m.process.Stop()
		m.process = nil
	}

	// Reset state
	m.stdio = []IoLine{}
	m.messages = nil
	m.extra = nil
	m.fs = nil

	// Create new process
	allArgs := []string{}
	allArgs = append(allArgs, m.config.Mode)
	allArgs = append(allArgs, m.config.Package)
	if m.config.Mode == "run" {
		allArgs = append(allArgs, m.config.Args...)
	}
	process, err := NewCommandProcess("go", allArgs...)

	if err != nil {
		m.systemErrors = append(m.systemErrors, err)
		return m, Error.Cmd()
	}
	m.process = process

	// Record start time
	start := time.Now()
	m.start = &start
	m.end = nil

	// Start the process
	if err := m.process.Start(); err != nil {
		m.systemErrors = append(m.systemErrors, err)
		return m, Error.Cmd()
	}

	return m, tea.Batch(status, m.checkCmdStatus(), m.checkStdout(), m.checkStderr(), updateViewport)
}

type stdoutMsg struct {
	line string
	seq  int
}

type stderrMsg struct {
	line string
	seq  int
}

func (m Model) checkStdout() tea.Cmd {
	if m.process == nil {
		return nil
	}

	return func() tea.Msg {
		select {
		case msg, ok := <-m.process.stdoutChan:
			if !ok {
				return nil
			}
			return msg
		case <-time.After(10 * time.Millisecond):
			return nil
		}
	}
}

func (m Model) checkStderr() tea.Cmd {
	if m.process == nil {
		return nil
	}

	return func() tea.Msg {
		select {
		case msg, ok := <-m.process.stderrChan:
			if !ok {
				return nil
			}
			return msg
		case <-time.After(10 * time.Millisecond):
			return nil
		}
	}
}

type processState struct {
	state *os.ProcessState
}

func (m Model) checkCmdStatus() tea.Cmd {
	if m.process == nil {
		return nil
	}

	return func() tea.Msg {
		return processState{m.process.CheckStatus()}
	}
}