package main

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

var conversation []api.Message
var ollama *api.Client

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

func roleToInt(role string) MessageType {
	if role == "system" {
		return MT_SYSTEM
	} else if role == "assistant" {
		return MT_ASSISTANT
	} else if role == "user" {
		return MT_USER
	} else {
		log.Fatalf("Invalid role type %s\n", role)
		return 0
	}
}

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 ""
	}
}

func loadMessageFromDb() []api.Message {
	dbMessages := GetMesages()
	var chatMessages []api.Message
	for _, msg := range dbMessages {
		message := api.Message{
			Role:    convertRole(msg.Type),
			Content: msg.Content,
		}
		chatMessages = append(chatMessages, message)
		fmt.Printf("%s: %s\n", colorRole(roleToInt(message.Role)), message.Content)
	}
	return chatMessages
}

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 onUserInput(input string) {
	if strings.HasPrefix(input, "/") {
		executeCommand(input)
		return
	}
	err := SaveMessage(MT_USER, input)
	if err != nil {
		log.Fatal(err)
	}

	conversation = append(conversation, api.Message{
		Role:    convertRole(MT_USER),
		Content: input,
	})

	ctx := context.Background()
	req := &api.ChatRequest{
		Model:    "llama3.2:1b",
		Messages: conversation,
	}

	fullResponse := ""
	respFunc := func(resp api.ChatResponse) error {
		fullResponse = fullResponse + resp.Message.Content
		fmt.Print(resp.Message.Content)
		if !resp.Done {
			return nil
		}
		return SaveMessage(MT_ASSISTANT, fullResponse)
	}

	fmt.Print(colorRole(MT_ASSISTANT), ": ")
	err = ollama.Chat(ctx, req, respFunc)
	if err != nil {
		log.Fatal(err)
	}
	conversation = append(conversation, api.Message{
		Role:    convertRole(MT_ASSISTANT),
		Content: fullResponse,
	})
	fmt.Println()
}

func main() {
	var err error
	if len(os.Args) <= 1 {
		log.Fatal("Missing command line parameter")
	}

	ollama, err = api.ClientFromEnvironment()
	if err != nil {
		log.Fatal(err)
	}

	dbName := os.Args[1]
	OpenDb(dbName)
	conversation = loadMessageFromDb()

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

	runner.Run()
}