package entity

import (
	"context"
	"gorm.io/gorm"
	"yx-dataset-server/app/schema"
)

// GetDatasetDB 获取Dataset存储
func GetDatasetDB(ctx context.Context, defDB *gorm.DB) *gorm.DB {
	return getDBWithModel(ctx, defDB, Dataset{})
}

// SchemaDataset Dataset对象
type SchemaDataset schema.Dataset

// ToDataset 转换为Dataset实体
func (a SchemaDataset) ToDataset() *Dataset {
	item := &Dataset{
		RecordID:  a.RecordID,
		Name:      a.Name,
		OrgId:     a.OrgId,
		FileCount: a.FileCount,
		Status:    a.Status,
		CreatorID: a.CreatorId,
		RagDataId: a.RagDataId,
		Sequence:  &a.Sequence,
	}
	return item
}

// ToSchemaDataset 转换为Dataset对象
func (a Dataset) ToSchemaDataset() *schema.Dataset {
	item := &schema.Dataset{
		RecordID:  a.RecordID,
		Name:      a.Name,
		OrgId:     a.OrgId,
		FileCount: a.FileCount,
		Status:    a.Status,
		CreatorId: a.CreatorID,
		RagDataId: a.RagDataId,
		CreatedAt: a.CreatedAt,
		Sequence:  *a.Sequence,
	}
	return item
}

// Dataset Dataset实体
type Dataset struct {
	gorm.Model
	RecordID  string `gorm:"column:record_id;size:32;index;"`
	Name      string `gorm:"column:name;size:30;index;"`
	OrgId     string `gorm:"column:org_id;size:32;index;"`
	FileCount int    `gorm:"column:file_count;"`
	Status    bool   `gorm:"column:status;"`
	CreatorID string `gorm:"column:creator_id;size:32;"`
	RagDataId string `gorm:"column:rag_data_id;size:32;"`
	Sequence  *int   `gorm:"column:sequence;auto_increment;"`
}

// Datasets Dataset列表
type Datasets []*Dataset

// ToSchemaDatasets 转换为Dataset对象列表
func (a Datasets) ToSchemaDatasets() []*schema.Dataset {
	list := make([]*schema.Dataset, len(a))
	for i, item := range a {
		list[i] = item.ToSchemaDataset()
	}
	return list
}