dynamo-browse/internal/common/ui/commandctrl/commandctrl.go

package commandctrl

import (
	"bytes"
	"context"
	"fmt"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/pkg/errors"
	"log"
	"os"
	"path/filepath"
	"strings"
	"ucl.lmika.dev/ucl"
	"ucl.lmika.dev/ucl/builtins"

	"github.com/lmika/shellwords"
	"lmika.dev/cmd/dynamo-browse/internal/common/ui/events"
)

const commandsCategory = "commands"

type cmdMessage struct {
	cmd string
}

type CommandController struct {
	uclInst            *ucl.Inst
	historyProvider    IterProvider
	commandList        *CommandList
	lookupExtensions   []CommandLookupExtension
	completionProvider CommandCompletionProvider
	uiStateProvider    UIStateProvider
	cmdChan            chan cmdMessage
	msgChan            chan tea.Msg
	interactive        bool
}

func NewCommandController(historyProvider IterProvider, pkgs ...CommandPack) (*CommandController, error) {
	cc := &CommandController{
		historyProvider:  historyProvider,
		commandList:      nil,
		lookupExtensions: nil,
		cmdChan:          make(chan cmdMessage),
		msgChan:          make(chan tea.Msg),
		interactive:      true,
	}

	options := []ucl.InstOption{
		ucl.WithOut(ucl.LineHandler(cc.printLine)),
		ucl.WithModule(builtins.CSV(nil)),
		ucl.WithModule(builtins.FS(nil)),
		ucl.WithModule(builtins.Log(nil)),
		ucl.WithModule(builtins.Itrs()),
		ucl.WithModule(builtins.OS()),
		ucl.WithModule(builtins.Strs()),
		ucl.WithModule(builtins.Lists()),
		ucl.WithModule(builtins.Time()),
	}
	for _, pkg := range pkgs {
		options = append(options, pkg.InstOptions()...)
	}

	cc.uclInst = ucl.New(options...)

	for _, pkg := range pkgs {
		pkg.ConfigureUCL(cc.uclInst)
	}

	execCtx := execContext{ctrl: cc}
	ctx := context.WithValue(context.Background(), commandCtlKey, &execCtx)
	for _, pkg := range pkgs {
		if err := pkg.RunPrelude(ctx, cc.uclInst); err != nil {
			return nil, err
		}
	}

	go cc.cmdLooper()

	return cc, nil
}

func (c *CommandController) AddCommands(ctx *CommandList) {
	ctx.parent = c.commandList
	c.commandList = ctx
}

func (c *CommandController) StartMessageSender(msgSender func(tea.Msg)) {
	for msg := range c.msgChan {
		msgSender(msg)
	}
}

func (c *CommandController) SetUIStateProvider(provider UIStateProvider) {
	c.uiStateProvider = provider
}

func (c *CommandController) AddCommandLookupExtension(ext CommandLookupExtension) {
	c.lookupExtensions = append(c.lookupExtensions, ext)
}

func (c *CommandController) SetCommandCompletionProvider(provider CommandCompletionProvider) {
	c.completionProvider = provider
}

func (c *CommandController) Prompt() tea.Msg {
	return events.PromptForInputMsg{
		Prompt:  ":",
		History: c.historyProvider.Iter(context.Background(), commandsCategory),
		OnDone: func(value string) tea.Msg {
			return c.Execute(value)
		},
		// TEMP
		OnTabComplete: func(value string) (string, bool) {
			if c.completionProvider == nil {
				return "", false
			}

			if strings.HasPrefix(value, "sa ") || strings.HasPrefix(value, "da ") {
				tokens := shellwords.Split(strings.TrimSpace(value))
				lastToken := tokens[len(tokens)-1]

				options := c.completionProvider.AttributesWithPrefix(lastToken)
				if len(options) == 1 {
					return value[:len(value)-len(lastToken)] + options[0], true
				}
			}
			return "", false
		},
		// END TEMP
	}
}

func (c *CommandController) Execute(commandInput string) tea.Msg {
	return c.execute(ExecContext{FromFile: false}, commandInput)
}

func (c *CommandController) execute(ctx ExecContext, commandInput string) tea.Msg {
	input := strings.TrimSpace(commandInput)
	if input == "" {
		return nil
	}

	select {
	case c.cmdChan <- cmdMessage{cmd: input}:
		// good
	default:
		return events.Error(errors.New("command currently running"))
	}

	return nil
}

func (c *CommandController) ExecuteAndWait(ctx context.Context, commandInput string) (any, error) {
	return c.uclInst.EvalString(ctx, commandInput)
}

func (c *CommandController) Invoke(invokable ucl.Invokable, args []any) (msg tea.Msg) {
	execCtx := execContext{ctrl: c}
	ctx := context.WithValue(context.Background(), commandCtlKey, &execCtx)

	res, err := invokable.Invoke(ctx, args)
	if err != nil {
		msg = events.Error(err)
	} else if res != nil {
		msg = events.StatusMsg(fmt.Sprint(res))
	}
	if execCtx.requestRefresh {
		c.postMessage(events.ResultSetUpdated{})
	}

	return msg
}

func (c *CommandController) cmdLooper() {
	execCtx := execContext{ctrl: c}
	ctx := context.WithValue(context.Background(), commandCtlKey, &execCtx)

	for {
		select {
		case cmdChan := <-c.cmdChan:
			res, err := c.ExecuteAndWait(ctx, cmdChan.cmd)
			if err != nil {
				c.postMessage(events.Error(err))
			} else if res != nil {
				c.postMessage(events.StatusMsg(fmt.Sprint(res)))
			}
			if execCtx.requestRefresh {
				c.postMessage(events.ResultSetUpdated{})
			}
		}
	}
}

func (c *CommandController) Alias(commandName string) Command {
	return func(ctx ExecContext, args ucl.CallArgs) tea.Msg {
		command := c.lookupCommand(commandName)
		if command == nil {
			return events.Error(errors.New("no such command: " + commandName))
		}

		return command(ctx, args)
	}
}

func (c *CommandController) lookupCommand(name string) Command {
	for ctx := c.commandList; ctx != nil; ctx = ctx.parent {
		if cmd, ok := ctx.Commands[name]; ok {
			return cmd
		}
	}

	for _, exts := range c.lookupExtensions {
		if cmd := exts.LookupCommand(name); cmd != nil {
			return cmd
		}
	}
	return nil
}

func (c *CommandController) LoadExtensions(ctx context.Context, baseDirs []string) error {
	log.Printf("loading extensions: %v", baseDirs)
	for _, baseDir := range baseDirs {
		baseDir = os.ExpandEnv(baseDir)
		descendIntoSubDirs := !strings.HasSuffix(baseDir, ".")

		if stat, err := os.Stat(baseDir); err != nil {
			if os.IsNotExist(err) {
				continue
			}
			return err
		} else if !stat.IsDir() {
			continue
		}

		log.Printf("walking %v", baseDir)
		if err := filepath.Walk(baseDir, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				return err
			}
			if info.IsDir() {
				if !descendIntoSubDirs && path != baseDir {
					return filepath.SkipDir
				}
				return nil
			}
			if strings.HasSuffix(info.Name(), ".ucl") {
				if err := c.ExecuteFile(ctx, path); err != nil {
					log.Println(err)
				}
				log.Printf("loaded %v\n", path)
			}
			return nil
		}); err != nil {
			return err
		}
	}
	return nil
}

func (c *CommandController) ExecuteFile(ctx context.Context, filename string) error {
	oldInteractive := c.interactive
	c.interactive = false
	defer func() {
		c.interactive = oldInteractive
	}()

	baseFilename := filepath.Base(filename)

	execCtx := execContext{ctrl: c}
	ctx = context.WithValue(context.Background(), commandCtlKey, &execCtx)

	if rcFile, err := os.ReadFile(filename); err == nil {
		if err := c.executeFile(ctx, rcFile); err != nil {
			return errors.Wrapf(err, "error executing %v", baseFilename)
		}
	} else {
		return errors.Wrapf(err, "error loading %v", baseFilename)
	}
	return nil
}

func (c *CommandController) executeFile(ctx context.Context, file []byte) error {
	if _, err := c.uclInst.Eval(ctx, bytes.NewReader(file), ucl.WithSubEnv()); err != nil {
		return err
	}

	return nil
}

func (c *CommandController) Inst() *ucl.Inst {
	return c.uclInst
}

func (c *CommandController) cmdInvoker(ctx context.Context, name string, args ucl.CallArgs) (any, error) {
	command := c.lookupCommand(name)
	if command == nil {
		return nil, errors.New("no such command: " + name)
	}

	res := command(ExecContext{}, args)
	if errMsg, isErrMsg := res.(events.ErrorMsg); isErrMsg {
		return nil, errMsg
	}
	return teaMsgWrapper{res}, nil
}

func (c *CommandController) printLine(s string) {
	if c.msgChan == nil || !c.interactive {
		log.Println(s)
		return
	}

	select {
	case c.msgChan <- events.StatusMsg(s):
	default:
		log.Println(s)
	}
}

func (c *CommandController) postMessage(msg tea.Msg) {
	if c.msgChan == nil {
		return
	}

	c.msgChan <- msg
}

type teaMsgWrapper struct {
	msg tea.Msg
}