JCS-pub/coordinator/internal/repl/repl.go

package repl

import (
	"context"
	"strings"

	"github.com/c-bata/go-prompt"
	"github.com/spf13/cobra"
	"gitlink.org.cn/cloudream/common/pkgs/async"
	"gitlink.org.cn/cloudream/jcs-pub/coordinator/internal/db"
	"gitlink.org.cn/cloudream/jcs-pub/coordinator/internal/ticktock"
)

var (
	cmdCtxKey = &CommandContext{}
)

type ReplEventChan = async.UnboundChannel[ReplEvent]

type ReplEvent interface {
	IsReplEvent() bool
}

type ExitEvent struct {
	ReplEvent
}

type Repl struct {
	prompt *prompt.Prompt
	evtCh  *ReplEventChan
	db     *db.DB
	tktk   *ticktock.TickTock
}

func New(db *db.DB, ticktock *ticktock.TickTock) *Repl {
	r := &Repl{
		evtCh: async.NewUnboundChannel[ReplEvent](),
		db:    db,
		tktk:  ticktock,
	}
	r.prompt = prompt.New(
		r.executor,
		r.completer,
		prompt.OptionPrefix(">>> "),
		prompt.OptionTitle("JCS Coordinator REPL"),
		prompt.OptionSetExitCheckerOnInput(r.exitChecker),
	)
	return r
}

func (r *Repl) Start() *ReplEventChan {
	go func() {
		r.prompt.Run()
		r.evtCh.Send(ExitEvent{})
	}()
	return r.evtCh
}

func (r *Repl) completer(d prompt.Document) []prompt.Suggest {
	return nil
}

func (r *Repl) executor(input string) {
	fields := strings.Fields(input)
	if len(fields) == 0 {
		return
	}

	RootCmd.SetArgs(fields)
	RootCmd.ExecuteContext(context.WithValue(context.Background(), cmdCtxKey, &CommandContext{
		repl: r,
	}))
}

func (r *Repl) exitChecker(input string, breakline bool) bool {
	if breakline && input == "exit" {
		return true
	}

	return false
}

var RootCmd = &cobra.Command{}

type CommandContext struct {
	repl *Repl
}

func GetCmdCtx(cmd *cobra.Command) *CommandContext {
	return cmd.Context().Value(cmdCtxKey).(*CommandContext)
}