package dao

import (
	"context"
	"database/sql"
	"fmt"
	"rag/consts/public"
	"rag/model/entity"

	"gitea.com/red-future/common/db/gfdb"
)

var DatasetIndex = new(datasetIndexDao)

type datasetIndexDao struct{}

// Insert 插入数据集索引
func (d *datasetIndexDao) Insert(ctx context.Context, index *entity.DatasetIndex) (id int64, err error) {
	_, err = gfdb.DB(ctx).Model(ctx, public.TableNameDatasetIndex).Data(index).Insert()
	if err != nil {
		return
	}
	return 0, nil
}

// GetByDatasetId 根据数据集ID获取索引
func (d *datasetIndexDao) GetByDatasetId(ctx context.Context, datasetId int64) (result *entity.DatasetIndex, err error) {
	err = gfdb.DB(ctx).Model(ctx, public.TableNameDatasetIndex).Where(entity.DatasetIndexCol.DatasetId, datasetId).Scan(&result)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return
	}
	return result, nil
}

// IncVectorCount 增加或减少向量数量
func (d *datasetIndexDao) IncVectorCount(ctx context.Context, id int64, delta int64) (err error) {
	_, err = gfdb.DB(ctx).Model(ctx, public.TableNameDatasetIndex).
		Where(entity.DatasetIndexCol.Id, id).
		Increment(entity.DatasetIndexCol.VectorCount, delta)
	return
}

func (d *datasetIndexDao) InsertIndex(ctx context.Context, indexName string) (err error) {
	prefix, err := gfdb.GetTablePrefix(ctx)
	if err != nil {
		return
	}
	sqlStr := fmt.Sprintf(`
		CREATE INDEX IF NOT EXISTS %s
		ON %s
		USING ivfflat (vector vector_cosine_ops)
		WHERE vector IS NOT NULL;
	`, indexName, prefix+public.TableNameDocumentChunk)
	_, err = gfdb.DB(ctx).Exec(ctx, sqlStr)
	return
}