package vector_data

import (
	"ai-chat-service/pkg/config"
	"context"
	"time"

	"github.com/tencent/vectordatabase-sdk-go/tcvectordb"
)

type tencentChatRecordsData struct {
	config   *config.Config
	vectorDB *tcvectordb.Client
}

func newTencentChatRecordsData(config *config.Config) (IChatRecordsData, error) {
	option := &tcvectordb.ClientOption{
		Timeout:            time.Second * time.Duration(config.Vector.Tencent.Timeout),
		MaxIdldConnPerHost: config.Vector.Tencent.MaxIdleConnPerHost,
		IdleConnTimeout:    time.Second * time.Duration(config.Vector.Tencent.IdleConnTimeout),
		ReadConsistency:    tcvectordb.ReadConsistency(config.Vector.Tencent.ReadConsistency),
	}
	client, err := tcvectordb.NewClient(config.Vector.Tencent.Url, config.Vector.Tencent.Username, config.Vector.Tencent.Pwd, option)
	if err != nil {
		return nil, err
	}
	return &tencentChatRecordsData{
		config:   config,
		vectorDB: client,
	}, nil
}

func (data *tencentChatRecordsData) UpsertData(ctx context.Context, list []*ChatRecord) error {
	database := data.config.Vector.Tencent.Database
	collection := CHAT_RECORDS
	coll := data.vectorDB.Database(database).Collection(collection)
	documentList := make([]tcvectordb.Document, 0, len(list))
	for _, l := range list {
		doc := tcvectordb.Document{Id: l.ID}
		doc.Fields = make(map[string]tcvectordb.Field, len(l.KVs))
		for k, v := range l.KVs {
			doc.Fields[k] = tcvectordb.Field{Val: v}
		}
		documentList = append(documentList, doc)
	}
	_, err := coll.Upsert(ctx, documentList)
	return err
}

func (data *tencentChatRecordsData) QueryData(ctx context.Context, text map[string][]string) (id string, score float32, err error) {
	database := data.config.Vector.Tencent.Database
	collection := CHAT_RECORDS
	coll := data.vectorDB.Database(database).Collection(collection)
	result, err := coll.SearchByText(ctx, text, &tcvectordb.SearchDocumentParams{
		Params: &tcvectordb.SearchDocParams{Ef: 100},
		Limit:  1,
	})
	if err != nil {
		return "", 0, err
	}
	if len(result.Documents) > 0 && len(result.Documents[0]) > 0 {
		doc := result.Documents[0][0]
		return doc.Id, doc.Score, nil
	}
	return "", 0, nil
}