package main

import (
	"context"
	"flag"
	"fmt"
	"github.com/ollama/ollama/api"
	"log"
	"os"
)

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 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 sendPromptInput(input string, handler func(response api.ChatResponse) error) {
	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
		err := handler(resp)
		if err != nil {
			return err
		}
		if !resp.Done {
			return nil
		}
		return SaveMessage(MT_ASSISTANT, fullResponse)
	}

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

const usage = `Usage:
llamachat [OPTION]... [CONVERSATION]

Options:
    -w, --web    Access via a web page.

If CONVERSATION is left out, the conversation will be kept in-memory. If it is
present, the conversation will be stored in a file called CONVERSATION.chat.
Alternatively, it can also be passed the literal string '-' (a singel dash) to
ensure an in-memory database is used.
`

func main() {
	flag.Usage = func() { fmt.Fprintf(os.Stderr, "%s\n", usage) }

	var err error
	var useWeb bool
	db := ""
	flag.BoolVar(&useWeb, "web", false, "run in web mode")
	flag.BoolVar(&useWeb, "w", false, "run in web mode")

	flag.Parse()

	db = flag.Arg(0)
	if db == "-" {
		db = ""
	}
	OpenDb(db)
	conversation = loadMessageFromDb()

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

	if useWeb {
		runAsWeb()
	} else {
		runAsCommandLine()
	}
}