package pimport

import (
	"context"

	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/tidbcloud/core-svc/idl/pbgen/proto/core"
	"github.com/tidbcloud/dataflow-service/internal/db/entity"
	"github.com/tidbcloud/dataflow-service/internal/grpc/grpcli"
	"github.com/tidbcloud/dataflow-service/pkg/account/accountdb"
	"github.com/tidbcloud/dataflow-service/pkg/account/accounttypes"
	"go.uber.org/zap"
)

// ImportTask is the entity object to manipulate in model.
type ImportTask struct {
	entity.Import
	// Cluster stores the cluster information for the imported cluster.
	Cluster *core.Cluster
	// Tenant stores the tenant information for the imported cluster.
	Tenant             *accounttypes.Tenant
	LightningSourceURL string
}

// LoadImportClusterInfo loads the cluster information for this import task.
func LoadImportClusterInfo(ctx context.Context, imp *ImportTask, forceReload bool) error {
	if imp.Cluster != nil && !forceReload {
		return nil
	}
	cluster, bizErr := grpcli.GetClusterWithDataplaneByID(ctx, uint64(imp.ClusterID))
	if bizErr != nil {
		return errors.Annotate(bizErr, "get cluster info error")
	}
	imp.Cluster = cluster
	return nil
}

// BatchLoadImportClusterInfo loads a list of import tasks' cluster informations in a single API call.
func BatchLoadImportClusterInfo(ctx context.Context, impTasks []*ImportTask) ([]*ImportTask, error) {
	if len(impTasks) <= 0 {
		// avoid listing clusters with no cluster ID specified, in this case, it will iterate all the clusters
		return nil, nil
	}
	clusterIDs := make([]uint64, len(impTasks))
	for i, imp := range impTasks {
		clusterIDs[i] = uint64(imp.ClusterID)
	}
	listResp, err := grpcli.ListClustersById(ctx, clusterIDs)
	if err != nil {
		return nil, errors.Annotate(err, "list clusters by ID error")
	}
	clusterMap := listResp.GetClusters()
	missingClusterImportIDs := []uint64{}
	imps := []*ImportTask{}
	for _, imp := range impTasks {
		cluster, ok := clusterMap[uint64(imp.ClusterID)]
		if !ok {
			missingClusterImportIDs = append(missingClusterImportIDs, uint64(imp.ID))
			continue
		}
		imp.Cluster = cluster
		imps = append(imps, imp)
	}
	if len(missingClusterImportIDs) > 0 {
		return imps, errors.Errorf("cannot find cluster information for some import: %v", missingClusterImportIDs)
	}
	return imps, nil
}

// BatchLoadImportTenantInfo loads a list of import tasks' tenant informations in a single API call.
// It depends on loading these imports' cluster information.
func BatchLoadImportTenantInfo(ctx context.Context, impTasks []*ImportTask) ([]*ImportTask, error) {
	if len(impTasks) <= 0 {
		// avoid listing clusters with no cluster ID specified, in this case, it will iterate all the clusters
		return nil, nil
	}
	var err error
	if impTasks, err = BatchLoadImportClusterInfo(ctx, impTasks); err != nil {
		log.Error("batch load cluster info error", zap.Error(err))
	}
	if len(impTasks) <= 0 {
		return nil, nil
	}
	tenantIDs := make([]uint64, len(impTasks))
	for i, imp := range impTasks {
		tenantIDs[i] = imp.Cluster.GetTenantId()
	}
	tenantMap, err := accountdb.ListOrgMapByIDs(tenantIDs)
	if err != nil {
		return nil, errors.Annotate(err, "list tenant info map by ids failed")
	}
	missingTenantImportIDs := []uint64{}
	imps := []*ImportTask{}
	for _, imp := range impTasks {
		tenant, ok := tenantMap[imp.Cluster.GetTenantId()]
		if !ok {
			missingTenantImportIDs = append(missingTenantImportIDs, uint64(imp.ID))
			continue
		}
		imp.Tenant = &tenant
		imps = append(imps, imp)
	}
	if len(missingTenantImportIDs) > 0 {
		return imps, errors.Errorf("cannot find tenant information for some import IDs: %v", missingTenantImportIDs)
	}
	return imps, nil
}