package mdb

import (
	"context"
	"time"
	"ymgb/database"

	"go.mongodb.org/mongo-driver/v2/bson"
)

type AiChatMessage struct {
	Id      bson.ObjectID `bson:"_id"`
	ChatID  string        `bson:"chatId"`
	Role    string        `bson:"role"`
	Message string        `bson:"message"`
	Index   int           `bson:"index"`
}

func GetRPChatHistory(db *database.Context, chatId string) ([]AiChatMessage, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	col := database.GetMongoCollection(db, "rp_chat_messages")
	cursor, err := col.Find(ctx, bson.M{"chatId": chatId})
	if err != nil {
		return nil, err
	}
	result := make([]AiChatMessage, 0)
	err = cursor.All(ctx, &result)
	return result, err
}
func UpdateRPChatHistory(db *database.Context, chatId, role, message string, index int) error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	col := database.GetMongoCollection(db, "rp_chat_messages")
	_, err := col.InsertOne(ctx, AiChatMessage{
		bson.NewObjectID(),
		chatId,
		role,
		message, index,
	})
	return err
}
func GetRPChatHistorySize(db *database.Context, chatId string) (int64, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	col := database.GetMongoCollection(db, "rp_chat_messages")
	return col.CountDocuments(ctx, bson.M{"chatId": chatId})
}
func DeleteRPChatEntry(db *database.Context, entry AiChatMessage) error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	col := database.GetMongoCollection(db, "rp_chat_messages")
	_, err := col.DeleteOne(ctx, bson.M{"chatId": entry.ChatID})
	return err
}