package controller

import (
	"context"

	"rag-local/kb/model/dto"
	"rag-local/kb/service"
)

type kgEntity struct{}

var KgEntity = &kgEntity{}

func (c *kgEntity) List(ctx context.Context, req *dto.ListKgEntityReq) (*dto.ListKgEntityRes, error) {
	list, total, err := service.KgEntityService.List(ctx, req.DatasetId, req.Page, req.PageSize)
	if err != nil {
		return nil, err
	}
	return &dto.ListKgEntityRes{
		List:     list,
		Total:    total,
		Page:     req.Page,
		PageSize: req.PageSize,
	}, nil
}

func (c *kgEntity) Sources(ctx context.Context, req *dto.SourcesKgEntityReq) (*dto.SourcesKgEntityRes, error) {
	list, err := service.KgEntityService.Sources(ctx, req.DatasetId)
	if err != nil {
		return nil, err
	}
	out := make([]*dto.KgEntitySourceItem, 0, len(list))
	for _, s := range list {
		out = append(out, &dto.KgEntitySourceItem{Name: s.Name, Files: s.Files})
	}
	return &dto.SourcesKgEntityRes{List: out}, nil
}