package tidbtools

import (
	"context"

	importintong "github.com/tidbcloud/dataflow-service-ng/idl/pbgen/proto/import_into/v1beta2"
	"github.com/tidbcloud/dataflow-service/idl/pbgen/proto/dataflow"
	"google.golang.org/grpc/status"
)

func (s *svc) ListImports(ctx context.Context, orgID, projectID, clusterID uint64, page, pageSize int32) (*dataflow.ListImportResp, error) {
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	var err error
	if ctx, err = withDataflowOpsUser(ctx, orgID, projectID); err != nil {
		return nil, err
	}
	resp, err := s.dataflow.Import.List(ctx, &dataflow.ListImportReq{
		OrgId:     orgID,
		ProjectId: projectID,
		ClusterId: clusterID,
		Page:      page,
		PageSize:  pageSize,
	})
	if err != nil {
		if st, ok := status.FromError(err); ok && st.Code() == codeImportClusterNotFound {
			return &dataflow.ListImportResp{}, nil
		}
		return nil, err
	}
	return resp, nil
}

// ----- Premium ImportInto methods (dataflow-service-ng import_into v1beta2) -----

func (s *svc) ListImportIntoTasks(ctx context.Context, tidbID string, pageSize int32) (*importintong.ListImportIntoTasksResponse, error) {
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	return s.dataflow.ImportInto.ListImportIntoTasks(ctx, &importintong.ListImportIntoTasksRequest{
		TidbId:   tidbID,
		PageSize: pageSize,
	})
}

func (s *svc) GetImportIntoTask(ctx context.Context, tidbID, taskID string) (*importintong.ImportIntoTask, error) {
	ctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	return s.dataflow.ImportInto.GetImportIntoTask(ctx, &importintong.GetImportIntoTaskRequest{
		TidbId:   tidbID,
		ImportId: taskID,
	})
}