package main

import (
	"fmt"
	"github.com/elk-language/go-prompt"
	"github.com/fatih/color"
	"github.com/ollama/ollama/api"
	"log"
	"strings"
)

func onUserInput(input string) {
	if strings.HasPrefix(input, "/") {
		executeCommand(input)
		return
	}
	fmt.Print(colorRole(MT_ASSISTANT), ": ")
	sendPromptInput(input, func(r api.ChatResponse) error {
		_, err := fmt.Print(r.Message.Content)
		return err
	})
	fmt.Println()
}

func runAsCommandLine() {
	for _, msg := range conversation {
		fmt.Printf("%s: %s\n", colorRole(roleToInt(msg.Role)), msg.Content)
	}

	runner := prompt.New(
		onUserInput,
		prompt.WithTitle("llamachat"),
		prompt.WithPrefix("user: "),
		prompt.WithCompleter(Completer),
	)

	runner.Run()
}

func executeCommand(cli string) {
	args := strings.Split(cli, " ")
	for _, cmd := range Commands {
		if cmd.Name == args[0] {
			cmd.Action(args)
			return
		}
	}
	fmt.Println("Unknown command")
}

func colorRole(role MessageType) string {
	if role == MT_SYSTEM {
		return color.RedString("system")
	} else if role == MT_ASSISTANT {
		return color.GreenString("assistant")
	} else if role == MT_USER {
		return color.BlueString("user")
	} else {
		log.Fatalf("Invalid role type %d\n", role)
		return ""
	}
}