package plugins

import (
	"fmt"
	"kurumibot/database/mdb"
	"kurumibot/database/psql"
	"kurumibot/database/red"
	"kurumibot/laniakea"
	"kurumibot/utils/ai"
	"os"
	"strconv"
	"strings"

	"github.com/google/uuid"
)

func RegisterTestRP(bot *laniakea.Bot) {
	rp := laniakea.NewPlugin("RP")
	rp = rp.Command(selectWaifu, "rpwaifu", "рпвайфу")
	rp = rp.Payload(selectWaifu, "rp.selwaifu")
	rp = rp.Command(newChat, "newchat")
	rp = rp.Command(generate, "g", "gen", "г")

	bot.AddPlugins(rp.Build())
}

func selectWaifu(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
	waifuId, err := strconv.Atoi(ctx.Args[0])
	if err != nil {
		ctx.Error(err)
		return
	}
	err = red.RPSetSelectedWaifu(db, ctx.FromID, waifuId)
	if err != nil {
		ctx.Error(err)
		return
	}

	ctx.Answer(fmt.Sprintf("Была выбрана вайфу %d", waifuId))
}

func newChat(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
	waifuId := red.RPGetSelectedWaifu(db, ctx.FromID)
	if waifuId == 0 {
		ctx.Answer("Не выбрана вайфу")
		return
	}
	chatId := uuid.New()
	err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId.String())
	if err != nil {
		ctx.Error(err)
		return
	}

	prompt := strings.Join(ctx.Args[1:], " ")
	err = red.RPSetChatPrompt(db, ctx.FromID, waifuId, prompt)
	if err != nil {
		ctx.Error(err)
		return
	}
	err = red.RPSetCounter(db, ctx.FromID, waifuId, 0)
	if err != nil {
		ctx.Error(err)
		return
	}
	ctx.Answer("Был создан новый чат.")
}

func generate(ctx *laniakea.MsgContext, db *laniakea.DatabaseContext) {
	waifuId := red.RPGetSelectedWaifu(db, ctx.FromID)
	if waifuId == 0 {
		ctx.Answer("Не выбрана вайфу")
		return
	}
	chatId := red.RPGetChatId(db, ctx.FromID, waifuId)
	if chatId == "" {
		chatId = uuid.New().String()
		err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId)
		if err != nil {
			ctx.Error(err)
			return
		}
	}
	waifu, err := psql.GetWaifuById(waifuId)
	if err != nil {
		ctx.Error(err)
		return
	}
	systemPrompt := ai.Message{
		Role: "system",
		Content: fmt.Sprintf(
			"%s %s %s",
			ai.FormatPrompt(ai.PreHistoryPrompt, waifu.Name, ctx.Msg.From.FirstName),
			fmt.Sprintf("Вот краткое описание твоего персонажа: %s", waifu.RpPrompt),
			red.RPGetChatPrompt(db, ctx.FromID, waifuId),
		),
	}

	history, err := mdb.GetChatHistory(db, chatId)
	if err != nil {
		ctx.Error(err)
		return
	}
	messages := make([]ai.Message, 0)
	for _, m := range history {
		messages = append(messages, ai.Message{
			Role:    m.Role,
			Content: m.Message,
		})
	}

	api := ai.NewOpenAIAPI(ai.CosmoRPUrl, os.Getenv("PAWAN_KEY"), "cosmorp-2.5")
	userMessage := strings.Join(ctx.Args, " ")
	messages = append(messages, ai.Message{
		Role: "system", Content: ai.FormatPrompt(ai.PostHistoryPrompt, waifu.Name, ctx.Msg.From.FirstName),
	}, ai.Message{
		Role: "user", Content: userMessage,
	})
	err = mdb.UpdateChatHistory(db, chatId, "user", userMessage)
	if err != nil {
		ctx.Error(err)
		return
	}
	m := ctx.Answer("Генерация запущена...")
	res, err := api.CreateCompletion(ai.CreateCompletionReq{
		Messages:    append([]ai.Message{systemPrompt}, messages...),
		Verbosity:   "low",
		Temperature: 0.7,
	})
	if err != nil {
		ctx.Error(err)
		return
	}
	response := make([]string, 0)
	for _, choice := range res.Choices {
		m := choice.Message
		response = append(response, m.Content)
		err = mdb.UpdateChatHistory(db, chatId, m.Role, m.Content)
	}
	m.Delete()
	ctx.Answer(laniakea.EscapeMarkdown(strings.Join(response, "\n")))

	counter := red.RPGetCounter(db, ctx.FromID, waifuId) + 1
	if counter == 5 {
		res, err = api.CreateCompletion(ai.CreateCompletionReq{
			Messages: append(messages, ai.Message{
				Role:    "user",
				Content: "Сделай краткую выжимку нашего диалога",
			}),
		})
		if err != nil {
			ctx.Error(err)
		}
		chatId = uuid.New().String()
		err := red.RPSetChatId(db, ctx.FromID, waifuId, chatId)
		if err != nil {
			ctx.Error(err)
			return
		}
		err = mdb.UpdateChatHistory(db, chatId, "assistant", res.Choices[0].Message.Content)
		err = red.RPSetCounter(db, ctx.FromID, waifuId, 0)
		if err != nil {
			ctx.Error(err)
			return
		}
	} else {
		err = red.RPSetCounter(db, ctx.FromID, waifuId, counter)
		if err != nil {
			ctx.Error(err)
		}
	}
}