package main

import (
	"Kevin/pkg/llama"
	"Kevin/pkg/structs"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/bwmarrin/discordgo"
	"github.com/spf13/viper"
)

var (
	messages []structs.Message = []structs.Message{}
	brain    *llama.LLama
	config   *viper.Viper
)

func main() {
	config = viper.New()
	config.SetConfigType("toml")
	config.SetConfigName("kevin")
	config.AddConfigPath(".")
	config.SetDefault("prefix", "!")
	config.SetDefault("model", "models/ggml-vicuna-7b-1.1-q4_0.bin")
	config.SetDefault("name", "Kevin")
	config.SetDefault("discord.token", "not-set")
	config.SetDefault("discord.chat", "not-set")
	config.SetDefault("discord.owner", "not-set")
	err := config.ReadInConfig()
	if err != nil {
		config.WriteConfigAs("kevin.toml")
		panic(fmt.Errorf("fatal error config file: %s", err))
	}

	if config.GetString("discord.token") == "not-set" {
		panic(fmt.Errorf("fatal error config file: %s", "discord token not set"))
	}

	brain, err = llama.New(config.GetString("model"), llama.SetContext(128), llama.SetParts(-1))
	if err != nil {
		panic(err)
	}

	discord, err := discordgo.New("Bot " + config.GetString("discord.token"))
	if err != nil {
		panic(err)
	}
	discord.AddHandler(messageCreate)
	discord.AddHandler(ready)
	discord.Identify.Intents = discordgo.MakeIntent(discordgo.IntentsGuildMessages)

	err = discord.Open()
	if err != nil {
		panic(err)
	}
	fmt.Println("Bot is now running. Press CTRL-C to exit.")
	sc := make(chan os.Signal, 1)
	signal.Notify(sc, syscall.SIGINT, syscall.SIGTERM, os.Interrupt)
	<-sc
	discord.Close()
	brain.Free()
}

func messageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
	if m.ChannelID != config.GetString("discord.chat") {
		return
	}
	if m.Author.ID == s.State.User.ID {
		return
	}
	if strings.HasPrefix(m.Content, config.GetString("prefix")) {
		if m.Author.ID == config.GetString("discord.owner") {
			if strings.HasPrefix(m.Content, "!clear") {
				messages = []structs.Message{}
				s.ChannelMessageSend(config.GetString("discord.chat"), "> Cleared messages")
			}
		}
		return
	}

	// logic
	s.UpdateStatusComplex(generateStatus(string(discordgo.StatusDoNotDisturb)))
	message := structs.NewMessage(m.Author.Username, strings.TrimSpace(m.Content))
	prompt := structs.GeneratePrompt(config, messages, *message)
	var seed int = int(time.Now().Unix())
	res, err := brain.Predict(prompt, llama.SetThreads(4), llama.SetTokens(128), llama.SetTopK(90), llama.SetTopP(0.86), llama.SetSeed(seed), llama.SetStopPrompt("###"))
	if err != nil {
		panic(err)
	}
	if res == "" {
		res = "> Something went wrong"
	}
	s.ChannelMessageSend(config.GetString("discord.chat"), res)
	message.Response = strings.TrimSpace(res)
	messages = append(messages, *message)
	s.UpdateStatusComplex(generateStatus(string(discordgo.StatusOnline)))
}

func ready(s *discordgo.Session, r *discordgo.Ready) {
	s.UpdateStatusComplex(generateStatus(string(discordgo.StatusOnline)))
	s.ChannelMessageSend(config.GetString("discord.chat"), "> Kevin is now running")
}

func generateStatus(status string) discordgo.UpdateStatusData {
	return discordgo.UpdateStatusData{
		AFK:    false,
		Status: status,
		Activities: []*discordgo.Activity{
			{
				Name: "chat",
				Type: discordgo.ActivityTypeListening,
			},
		},
	}
}