YaeMikoBot/database/red/rp_chats.go

package red

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"strings"
	"ymgb/database"
	"ymgb/database/psql"
	"ymgb/utils"

	"github.com/google/uuid"
	"github.com/redis/go-redis/v9"
)

var ctx = context.Background()

type RPRepository struct {
	client *redis.Client
	db     *database.Context
}

type RPChat struct {
	ID         uuid.UUID
	UserID     int
	WaifuID    int
	Prompt     string
	Counter    int
	ChatTokens int64

	SettingID    int
	Setting      *psql.RPSetting
	ScenariosIDs []int
	Scenarios    []psql.RPScenario
}

func NewRPRepository(db *database.Context) RPRepository {
	return RPRepository{db.Redis, db}
}

func (rep RPRepository) SetSelectedWaifu(userId, waifuId int) error {
	key := fmt.Sprintf("ai.chats.rp.%d", userId)
	return rep.client.Set(ctx, key, waifuId, 0).Err()
}
func (rep RPRepository) GetSelectedWaifu(userId int) int {
	key := fmt.Sprintf("ai.chats.rp.%d", userId)
	res := rep.client.Get(ctx, key)
	if res.Err() != nil {
		return 0
	}
	i, _ := res.Int()
	return i
}

func (rep RPRepository) SetChatId(userId, waifuId int, chatId string) error {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.chat", userId, waifuId)
	return rep.client.Set(context.Background(), key, chatId, 0).Err()
}
func (rep RPRepository) GetChatId(userId, waifuId int) string {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.chat", userId, waifuId)
	res := rep.client.Get(ctx, key)
	if res.Err() != nil {
		return ""
	}
	return res.Val()
}
func (rep RPRepository) GetOrCreateChatId(userId, waifuId int) (string, error) {
	chatId := rep.GetChatId(userId, waifuId)
	if chatId == "" {
		chatId = uuid.New().String()
	}
	err := rep.SetChatId(userId, waifuId, chatId)
	return chatId, err
}
func (rep RPRepository) GetChat(userId int) (RPChat, error) {
	var chat RPChat
	waifuId := rep.GetSelectedWaifu(userId)
	if waifuId == 0 {
		return chat, errors.New("no WaifuID")
	}

	chatId, err := rep.GetOrCreateChatId(userId, waifuId)
	if err != nil {
		return chat, err
	}

	chat.ID = uuid.MustParse(chatId)
	chat.UserID = userId
	chat.WaifuID = waifuId
	chat.Prompt = rep.GetChatPrompt(userId, waifuId)
	chat.Counter = rep.GetCounter(userId, waifuId)
	chat.ChatTokens = int64(rep.GetChatTokens(userId, waifuId))

	chat.SettingID = rep.GetChatSettingID(userId, waifuId)
	psqlRep := psql.NewRPRepository(rep.db)
	if chat.SettingID > 0 {
		setting, err := psqlRep.GetSetting(chat.SettingID)
		if err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				return chat, fmt.Errorf("setting %d not found", chat.SettingID)
			}
			return chat, err
		}
		chat.Setting = &setting
	}

	chat.ScenariosIDs = rep.GetChatScenariosIDs(userId, waifuId)
	chat.Scenarios = make([]psql.RPScenario, len(chat.ScenariosIDs))
	for i, id := range chat.ScenariosIDs {
		if id <= 0 {
			continue
		}
		chat.Scenarios[i], err = psqlRep.GetScenario(id)
		if err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				return chat, fmt.Errorf("scenario %d not found", id)
			}
			return chat, err
		}
	}
	return chat, nil
}

func (rep RPRepository) SaveChat(chat RPChat) error {
	waifuId := chat.WaifuID
	userId := chat.UserID
	var err error

	if err = rep.SetChatPrompt(userId, waifuId, chat.Prompt); err != nil {
		return err
	}
	if err = rep.SetCounter(userId, waifuId, chat.Counter); err != nil {
		return err
	}
	if err = rep.SetChatTokens(userId, waifuId, int(chat.ChatTokens)); err != nil {
		return err
	}

	var settingId int
	if chat.Setting != nil {
		settingId = chat.Setting.ID
	} else {
		settingId = chat.SettingID
	}
	if err = rep.SetChatSettingID(userId, waifuId, settingId); err != nil {
		return err
	}

	var ids []int
	if len(chat.Scenarios) > 0 {
		ids = utils.Map(chat.Scenarios, func(s psql.RPScenario) int {
			return s.ID
		})
	} else {
		ids = chat.ScenariosIDs
	}
	if err = rep.SetChatScenariosIDs(userId, waifuId, ids); err != nil {
		return err
	}
	return nil
}

func (rep RPRepository) SetChatPrompt(userId, waifuId int, prompt string) error {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.prompt", userId, waifuId)
	return rep.client.Set(context.Background(), key, prompt, 0).Err()
}
func (rep RPRepository) GetChatPrompt(userId, waifuId int) string {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.prompt", userId, waifuId)
	res := rep.client.Get(ctx, key)
	if res.Err() != nil {
		return ""
	}
	return res.Val()
}

func (rep RPRepository) SetCounter(userId, waifuId, counter int) error {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.counter", userId, waifuId)
	return rep.client.Set(ctx, key, counter, 0).Err()
}
func (rep RPRepository) GetCounter(userId, waifuId int) int {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.counter", userId, waifuId)
	res := rep.client.Get(ctx, key)
	if res.Err() != nil {
		return 0
	}
	i, _ := res.Int()
	return i
}

func (rep RPRepository) SetChatTokens(userId, waifuId, tokens int) error {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.tokens", userId, waifuId)
	return rep.client.Set(ctx, key, tokens, 0).Err()
}
func (rep RPRepository) GetChatTokens(userId, waifuId int) int {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.tokens", userId, waifuId)
	res := rep.client.Get(ctx, key)
	if res.Err() != nil {
		return 0
	}
	i, _ := res.Int()
	return i
}

func (rep RPRepository) SetChatSettingID(userId, waifuId, settingId int) error {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.setting_id", userId, waifuId)
	return rep.client.Set(ctx, key, settingId, 0).Err()
}
func (rep RPRepository) GetChatSettingID(userId, waifuId int) int {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.setting_id", userId, waifuId)
	res := rep.client.Get(ctx, key)
	if res.Err() != nil {
		return 0
	}
	i, _ := res.Int()
	return i
}

func (rep RPRepository) SetChatScenariosIDs(userId, waifuId int, scenarioIds []int) error {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.scenario_id", userId, waifuId)
	ids := strings.Join(utils.Map(scenarioIds, utils.AnyToString), ",")
	return rep.client.Set(ctx, key, ids, 0).Err()
}
func (rep RPRepository) GetChatScenariosIDs(userId, waifuId int) []int {
	key := fmt.Sprintf("ai.chats.rp.%d.%d.scenario_id", userId, waifuId)
	res := rep.client.Get(ctx, key)
	if res.Err() != nil {
		return []int{0}
	}
	ids := strings.Split(res.Val(), ",")
	return utils.Map(ids, utils.StringToInt)
}