package pimport

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

	infraapi "github.com/tidbcloud/central-sdk/infra-api"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"

	gardencorev1beta1 "github.com/gardener/gardener/pkg/apis/core/v1beta1"
	"github.com/pingcap/dbaas/pkg/capacity"
	"github.com/pingcap/dbaas/pkg/secret/encrypt"
	"github.com/pingcap/dbaas/pkg/secret/vault"
	"github.com/pingcap/dbaas/pkg/tidb/component"
	"github.com/pingcap/errcode"
	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/tidbcloud/central-sdk/config"
	"github.com/tidbcloud/central-sdk/db"
	"github.com/tidbcloud/central-sdk/db/dbutil"
	"github.com/tidbcloud/central-sdk/db/entity"
	"github.com/tidbcloud/central-sdk/db/tenant"
	central "github.com/tidbcloud/central-sdk/proto/output"
	"github.com/tidbcloud/central-sdk/runtime"
	"github.com/tidbcloud/central-sdk/runtime/deployer"
	"github.com/tidbcloud/central-sdk/runtime/deployer/gardener"
	"github.com/tidbcloud/central-sdk/runtime/deployer/k8s"
	runtimemanager "github.com/tidbcloud/central-sdk/runtime/manager"
	runtimeutil "github.com/tidbcloud/central-sdk/runtime/util"
	"go.uber.org/multierr"
	"go.uber.org/zap"
	"gorm.io/gorm"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/util/retry"
)

const (
	DatasourceS3  = 1
	DataSourceGCS = 2
)

func identity(imp *entity.Import) uint64 {
	id := uint64(imp.ID)
	// OldID is the id of the import table per-tenant
	// We can deprecated the OldID once migrated to new schema
	if imp.OldID != nil {
		id = *imp.OldID
	}

	return id
}

func name(imp *entity.Import) string {
	id := identity(imp)
	return "tidb-lightning-for-import-" + strconv.FormatUint(id, 10)
}

const validImportTimeout = time.Second * 5

// Model is created with NewModel
type Model struct {
	log       *zap.Logger
	runtimes  runtimemanager.Manager
	encryptor encrypt.Encryptor
	deployer  deployer.ClusterDeployer
	api       *infraapi.API
}

// NewModel returns a Model
func NewModel(log *zap.Logger, runtimes runtimemanager.Manager, infraAPI *infraapi.API, encryptor encrypt.Encryptor, getDeployer deployer.ClusterDeployer) Model {
	return Model{
		log:       log,
		runtimes:  runtimes,
		encryptor: encryptor,
		deployer:  getDeployer,
		api:       infraAPI,
	}
}

// UserContext attaches a User to a Context.
// This helps audit the User
type UserContext struct {
	Model
	tenantAndUser *tenant.User
	log           *zap.Logger
}

// User returns a UserContext
func (m Model) User(user *tenant.User) UserContext {
	return UserContext{m, user, m.log.With(
		zap.Uint64("tenantID", uint64(user.TenantID)),
		zap.Uint64("userID", uint64(user.ID)),
		zap.String("module", "import"),
	)}
}

func (m Model) TenantLog(tenantID uint64) *zap.Logger {
	return m.log.With(zap.Uint64("tenantID", tenantID), zap.String("model", "import"))
}

func (c UserContext) CheckOrgIDAndProjectID(orgID, projectID entity.ID) error {
	user := c.tenantAndUser
	if user.IsImpersonated {
		return nil
	}

	if orgID != user.TenantID {
		return errcode.NewForbiddenErr(errors.New("org_id parameter is inconsistent with the user's org ID"))
	}

	return dbutil.CheckProjectID(user.Projects, projectID)
}

func (c UserContext) Log() *zap.Logger {
	return c.log
}

func (c UserContext) ExitImport(req *central.ClusterImportId) (err error) {
	return db.Transaction(func(tx *gorm.DB) error {
		importID := entity.ID(req.Id)
		orgID := entity.ID(req.OrgId)
		projectID := entity.ID(req.ProjectId)
		clusterID := entity.ID(req.ClusterId)

		imp, err := dbutil.GetImportByID(tx, orgID, projectID, clusterID, importID)
		if err != nil {
			return err
		}

		if !isEndStatus(imp.Status) {
			return errors.Errorf("can't exit at status: %s", imp.Status)
		}
		imp.Cluster.ImportID = nil
		err = imp.Cluster.Update(tx, entity.ClusterDBSchema.ImportID)
		_ = dbutil.AddAccessLogNow(tx, c.tenantAndUser.User.TenantID, c.tenantAndUser.Email, fmt.Sprintf("exit import for cluster:  %d, import id: %d", clusterID, imp.ID))
		return errors.AddStack(err)
	})
}

// GetRoleARN get the role ARN we will used to run lightning
func (c UserContext) GetRoleARN(ctx context.Context, cluster *entity.Cluster) (string, error) {
	var (
		region *entity.Region
		err    error
	)
	err = db.Transaction(func(tx *gorm.DB) error {
		region, err = dbutil.GetClusterRegion(tx, cluster)
		return err
	})
	if err != nil {
		return "", err
	}

	var roleARN string
	if c.api.InfraAPIConfig.EnableNetwork {
		network, err := c.api.GetNetwork(ctx, cluster.Provider.Name.String(), cluster.ProjectID.Uint64(), region.Name)
		if err != nil {
			return "", err
		}
		if network.Status.Import.RoleARN == "" {
			return "", errors.New("import role is not ready")
		}
		roleARN = network.Status.Import.RoleARN
	} else {
		shootName := runtimeutil.ShootName(cluster.ProjectID.Uint64(), cluster.Provider.Name, region.Name)
		g := c.gardenEnv()
		roleARN, err = g.GetImportIamIdentity(g.ProjectName, shootName, c.tenantAndUser.TenantID.Uint64(), cluster.Provider.Name)
		if err != nil {
			return "", err
		}
	}

	return roleARN, nil
}

// CancelImport cancel the import
// we update the status to be canceling and wait it reach canceled if waitCanceled is true.
func (c UserContext) CancelImport(ctx context.Context, req *central.ClusterImportId, waitCanceled bool) (err error) {
	importID := entity.ID(req.Id)
	orgID := entity.ID(req.OrgId)
	projectID := entity.ID(req.ProjectId)
	clusterID := entity.ID(req.ClusterId)

	var imp = &entity.Import{}
	err = db.Transaction(func(tx *gorm.DB) error {
		imp, err = dbutil.GetImportByID(tx, orgID, projectID, clusterID, importID)
		if err != nil {
			return err
		}

		switch imp.Status {
		case entity.ImportStatusPreparing:
		case entity.ImportStatusImporting:
		case entity.ImportStatusCanceling:
		case entity.ImportStatusFailed:
			return errors.New("import already failed")
		case entity.ImportStatusCompleted:
			return errors.New("import already completed")
		case entity.ImportStatusCanceled:
			return nil
		default:
			return errors.Errorf("unknown status: %s", imp.Status)
		}

		// Handle cancel here.
		imp.Status = entity.ImportStatusCanceling
		err = imp.Update(tx, entity.ImportDBSchema.Status)
		if err != nil {
			return errors.AddStack(err)
		}

		return dbutil.AddAccessLogNow(tx, c.tenantAndUser.User.TenantID, c.tenantAndUser.Email, fmt.Sprintf("cancel import for cluster:  %d, import id: %d", clusterID, imp.ID))
	})
	if err != nil {
		return errors.AddStack(err)
	}

	if !waitCanceled {
		return nil
	}

	for {
		err := db.Transaction(func(tx *gorm.DB) error {
			return entity.NewImportQuerySet(tx).IDEq(importID).One(imp)
		})
		if err != nil {
			return errors.AddStack(err)
		}

		if isEndStatus(imp.Status) {
			return nil
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(time.Second):
		}
	}
}

func (c UserContext) GetImport(req *central.ClusterImportId) (*entity.Import, error) {
	var imp = &entity.Import{}
	var err error
	orgID := entity.ID(req.OrgId)
	projectID := entity.ID(req.ProjectId)
	clusterID := entity.ID(req.ClusterId)
	importID := entity.ID(req.Id)
	err = db.Transaction(func(tx *gorm.DB) error {
		imp, err = dbutil.GetImportByID(tx, orgID, projectID, clusterID, importID)
		if err != nil {
			return err
		}
		return nil
	})
	return imp, err
}

func (c UserContext) validImport(imp *entity.Import, cluster *entity.Cluster) (valid bool, reason string, err error) {
	ctx, cancel := context.WithTimeout(context.Background(), validImportTimeout)
	defer cancel()

	roleARN, err := c.GetRoleARN(ctx, cluster)
	if err != nil {
		return false, "", err
	}

	return validateReq(ctx, imp.SourceURL, imp.BucketRegion, roleARN)
}

func (c UserContext) CreateImport(orgID entity.ID, projectID entity.ID, imp *entity.Import) (id uint64, err error) {
	if imp.TargetPassword != "" {
		encryptedPassword, err := encrypt.EncryptBase64(c.encryptor, vault.Secret(imp.TargetPassword))
		if err != nil {
			return 0, err
		}

		imp.TargetPassword = encryptedPassword
	}

	var cluster entity.Cluster
	err = db.Transaction(func(tx *gorm.DB) error {
		cluster, err = dbutil.GetClusterByID(tx, orgID, projectID, imp.ClusterID)
		return err
	})
	if err != nil {
		return 0, errors.AddStack(err)
	}

	if !config.Global().ImportConfig.SkipBanMyBuckets {
		bucket := parseBucketName(imp.SourceURL)
		isMy, err := isMyBuckets(imp.BucketRegion, bucket)
		if err != nil {
			return 0, err
		}
		if isMy {
			return 0, errcode.NewInvalidInputErr(fmt.Errorf("bucket %s belong to forbidden account", bucket))
		}
	}

	if !config.Global().ImportConfig.SkipPreCheck {
		valid, reason, err := c.validImport(imp, &cluster)
		if err != nil {
			c.Log().Error("failed to validate import", zap.Reflect("import", imp), zap.Error(err))
		} else {
			if !valid {
				c.Log().Info("validate import result", zap.Bool("valid", valid), zap.String("reason", reason))
				return 0, errcode.NewInvalidInputErr(fmt.Errorf(reason))
			}
			c.log.Debug("pass validate import")
		}
	}

	err = db.Transaction(func(tx *gorm.DB) error {
		cluster, err := dbutil.GetClusterByID(tx, orgID, projectID, imp.ClusterID)
		if err != nil {
			return errors.AddStack(err)
		}

		if cluster.ImportID != nil {
			return errors.New("already importing")
		}

		if err := imp.Create(tx); err != nil {
			return errors.AddStack(err)
		}

		cluster.ImportID = new(uint64)
		*cluster.ImportID = uint64(imp.ID)
		if err := cluster.Update(tx, entity.ClusterDBSchema.ImportID); err != nil {
			return errors.AddStack(err)
		}

		return dbutil.AddAccessLogNow(tx, c.tenantAndUser.User.TenantID, c.tenantAndUser.Email, fmt.Sprintf("create import for cluster:  %d, import id: %d", cluster.ID, imp.ID))
	})

	if err != nil {
		return 0, err
	}

	return uint64(imp.ID), nil
}

// BgContext is meant for system worker/controller
type BgContext UserContext

// BgContext returns a BgContext
func (m Model) BgContext(tenantID entity.ID) (BgContext, error) {
	var systemUser entity.User
	err := db.Transaction(func(tx *gorm.DB) error {
		var err error
		systemUser, err = dbutil.SystemUser(tx, tenantID)
		return err
	})
	if err != nil {
		return BgContext{}, err
	}
	return BgContext{m, &tenant.User{TenantID: tenantID}, m.log.With(zap.Uint64("tenantID", tenantID.Uint64()), zap.Uint64("userID", uint64(systemUser.ID)), zap.String("module", "import"))}, nil
}

func (c BgContext) transToStatus(imp *entity.Import, status entity.ImportStatus, msg string) error {
	imp.Status = status
	var fields []entity.ImportDBSchemaField
	fields = append(fields, entity.ImportDBSchema.Status)

	if status == entity.ImportStatusImporting && imp.StartedAt == nil {
		t := time.Now()
		imp.StartedAt = &t
		fields = append(fields, entity.ImportDBSchema.StartedAt)
	}

	if isEndStatus(status) && imp.FinishedAt == nil {
		t := time.Now()
		imp.FinishedAt = &t
		fields = append(fields, entity.ImportDBSchema.FinishedAt)
	}

	if msg != "" {
		imp.Message = msg
		fields = append(fields, entity.ImportDBSchema.Message)
	}

	err := db.Transaction(func(tx *gorm.DB) error {
		return imp.Update(tx, fields...)
	})
	if err != nil {
		return errors.AddStack(err)
	}

	if isEndStatus(status) {
		return c.cleanUp(imp)
	}

	return nil
}

func isEndStatusV1(status entity.ImportStatus) bool {
	switch status {
	case entity.ImportStatus(cloudv1.ImportCompleted),
		entity.ImportStatus(cloudv1.ImportFailed):
		return true
	default:
		return false
	}
}

func isEndStatus(status entity.ImportStatus) bool {
	switch status {
	case entity.ImportStatusCompleted,
		entity.ImportStatusFailed,
		entity.ImportStatusCanceled:
		return true
	default:
		return false
	}
}

// delete worker created and resources created at k8s.
func (c BgContext) cleanUp(imp *entity.Import) error {
	if c.api.InfraAPIConfig.EnableImportWrite {
		err := c.api.DeleteImport(context.TODO(), uint64(imp.Cluster.ProjectID), uint64(imp.Cluster.ID))
		if err != nil {
			return err
		}

		imp.CleanUp = true
		err = db.Transaction(func(tx *gorm.DB) error {
			return imp.Update(tx, entity.ImportDBSchema.CleanUp)
		})
		if err != nil {
			return errors.AddStack(err)
		}

	} else {
		if err := c.cleanJob(imp); err != nil {
			// if cluster is deleted ignore failed to clean up
			// the shoot may not be available if we already remove the shoot.
			if imp.Cluster.Status.DeletedOrCleared() {
				c.log.Info("failed to cleanup job", zap.Error(err))
			} else {
				return err
			}
		}

		if err := c.cleanWorker(c.gardenEnv(), imp.Cluster, imp); err != nil {
			return err
		}
	}

	// currently we only fetch the imports not clean up to sync
	// but we cleanup also when in canceling status so make sure
	// update db status only when it's at the end status that we no longer need to sync the import.
	if isEndStatus(imp.Status) {
		imp.CleanUp = true
		err := db.Transaction(func(tx *gorm.DB) error {
			return imp.Update(tx, entity.ImportDBSchema.CleanUp)
		})
		if err != nil {
			return errors.AddStack(err)
		}
	}
	return nil
}

func (c BgContext) OnPreparing(imp *entity.Import) error {
	if c.api.InfraAPIConfig.EnableImportWrite {
		var password string
		if imp.TargetPassword != "" {
			secret, err := encrypt.DecryptBase64(c.encryptor, imp.TargetPassword)
			if err != nil {
				return errors.Annotatef(err, "failed to decrypt, import id: %d", imp.ID)
			}
			password = string(secret)
		}
		if imp.DataSourceType == DatasourceS3 {
			err := c.api.CreateAwsImportIgnoreExists(
				context.TODO(),
				uint64(imp.Cluster.TenantID),
				uint64(imp.Cluster.ProjectID),
				uint64(imp.ClusterID),
				imp.TargetUsername,
				password,
				imp.SourceURL,
				imp.BucketRegion,
				imp.Filter,
				cloudv1.DataFormat(imp.DataFormat),
			)
			if err != nil {
				return err
			}
		}
		if imp.DataSourceType == DataSourceGCS {
			err := c.api.CreateGcpImportIgnoreExists(
				context.TODO(),
				uint64(imp.Cluster.TenantID),
				uint64(imp.Cluster.ProjectID),
				uint64(imp.ClusterID),
				imp.TargetUsername,
				password,
				imp.SourceURL,
				imp.BucketRegion,
				imp.Filter,
				cloudv1.DataFormat(imp.DataFormat),
			)
			if err != nil {
				return err
			}
		}
	} else {
		g := c.gardenEnv()

		if err := c.ensureWorkerCreated(g, imp.Cluster, imp); err != nil {
			return err
		}

		if err := c.ensureJobCreated(g, imp); err != nil {
			return err
		}
	}

	return c.updateInfo(imp)
}

func (c BgContext) updateInfo(imp *entity.Import) error {
	var errs error
	if c.api.InfraAPIConfig.EnableImportRead {
		err := c.updateFromLightningV1(imp)
		if err != nil {
			errs = multierr.Append(errs, err)
		}

		if isEndStatusV1(imp.Status) {
			return errs
		}
	} else {
		err := c.updateFromLightning(imp)
		if err != nil {
			errs = multierr.Append(errs, err)
		}

		if isEndStatus(imp.Status) {
			return errs
		}

		err = c.updateFromJob(imp)
		if err != nil {
			errs = multierr.Append(errs, err)
		}
	}
	return errs
}

func convertStatus(status string) string {
	switch status {
	case string(cloudv1.ImportCompleted):
		return string(entity.ImportStatusCompleted)
	case string(cloudv1.ImportFailed):
		return string(entity.ImportStatusFailed)
	case string(cloudv1.ImportPrepare):
		return string(entity.ImportStatusPreparing)
	case string(cloudv1.ImportRunning):
		return string(entity.ImportStatusImporting)
	default:
		return string(entity.ImportStatusPreparing)
	}
}

func (c BgContext) updateFromLightningV1(imp *entity.Import) error {
	impcr, err := c.api.GetImport(context.TODO(), uint64(imp.Cluster.ProjectID), uint64(imp.Cluster.ID))
	if err != nil {
		return err
	}
	imp.TotalSize = &impcr.Status.TotalSize
	imp.TotalFiles = impcr.Status.TotalFiles
	imp.CurrentTables.Tables = nil
	for _, v := range impcr.Status.CurrentTables {
		imp.CurrentTables.Tables = append(imp.CurrentTables.Tables, &entity.CurrentTable{
			Name:             v.Name,
			Size:             v.Size,
			CompletedPercent: v.CompletedPercent,
		})
	}
	imp.CompletedTables = impcr.Status.CompletedTables
	imp.PendingTables = impcr.Status.PendingTables
	imp.CompletedPercent = impcr.Status.CompletedPercent
	imp.Status = entity.ImportStatus(convertStatus(string(impcr.Status.Phase)))
	imp.Message = impcr.Status.Message
	if len(imp.Message) > 255 {
		imp.Message = imp.Message[:255]
	}
	c.log.Debug("updating info from import object succeeded")
	err = db.Transaction(func(tx *gorm.DB) error {
		return imp.Update(tx,
			entity.ImportDBSchema.TotalSize,
			entity.ImportDBSchema.TotalFiles,
			entity.ImportDBSchema.CurrentTables,
			entity.ImportDBSchema.CompletedTables,
			entity.ImportDBSchema.PendingTables,
			entity.ImportDBSchema.CompletedPercent,
			entity.ImportDBSchema.Status,
			entity.ImportDBSchema.Message)
	})
	if err != nil {
		return errors.Annotate(err, "failed to update info to db")
	}
	return nil
}

func (c BgContext) updateFromLightning(imp *entity.Import) error {
	cluster := imp.Cluster
	shootCli, err := c.getShootCli(cluster)
	if err != nil {
		return err
	}

	ns := runtimeutil.TiDBNamespace(cluster.ID.Uint64()).String()

	oldStatus := imp.Status

	proxyGetter := newProxyProgressGetter(shootCli, ns, name(imp), "8289")
	err = updateLightningInfo(imp, proxyGetter)
	if err != nil {
		if strings.Contains(err.Error(), "the server is currently unable to handle the request") {
			c.log.Debug("service unready", zap.String("err", err.Error()))
			return nil
		}

		c.log.Error("failed to update info from proxy getter", zap.Error(err))
		if err != nil {
			return err
		}
	}
	c.log.Debug("update info from proxy getter success")

	err = db.Transaction(func(tx *gorm.DB) error {
		return imp.Update(tx,
			entity.ImportDBSchema.TotalSize,
			entity.ImportDBSchema.TotalFiles,
			entity.ImportDBSchema.CurrentTables,
			entity.ImportDBSchema.CompletedTables,
			entity.ImportDBSchema.PendingTables,
			entity.ImportDBSchema.CompletedPercent,
			entity.ImportDBSchema.Message)
	})
	if err != nil {
		return errors.Annotate(err, "failed to update info to db")
	}

	if imp.Status != oldStatus {
		return c.transToStatus(imp, imp.Status, "")
	}

	c.log.Debug("update current table", zap.Reflect("tables", imp.CurrentTables))

	return nil
}

func (c BgContext) updateFromJob(imp *entity.Import) error {
	cluster := imp.Cluster
	shootCli, err := c.getShootCli(cluster)
	if err != nil {
		return err
	}

	ns := runtimeutil.TiDBNamespace(cluster.ID.Uint64()).String()

	job, err := shootCli.BatchV1().Jobs(ns).Get(name(imp), metav1.GetOptions{})
	if err != nil {
		return errors.AddStack(err)
	}

	for _, cond := range job.Status.Conditions {
		if cond.Status != corev1.ConditionTrue {
			continue
		}

		if cond.Type == batchv1.JobFailed {
			return c.transToStatus(imp, entity.ImportStatusFailed, "")
		} else if cond.Type == batchv1.JobComplete {
			return c.transToStatus(imp, entity.ImportStatusCompleted, "")
		}
	}

	return nil
}

func (c BgContext) OnImporting(imp *entity.Import) error {
	return c.updateInfo(imp)
}

func (c BgContext) OnCanceling(imp *entity.Import) error {
	if c.api.InfraAPIConfig.EnableImportWrite {
		err := c.api.DeleteImport(context.TODO(), uint64(imp.Cluster.ProjectID), uint64(imp.Cluster.ID))
		if err != nil {
			return err
		}
	} else {
		err := c.cleanUp(imp)
		if err != nil {
			return err
		}

		ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
		defer cancel()

		err = c.waitLightningStop(ctx, imp)
		if err != nil {
			if err == context.DeadlineExceeded {
				return nil
			}
			return err
		}
	}
	return c.transToStatus(imp, entity.ImportStatusCanceled, "")
}

func (c BgContext) sync(imp *entity.Import) error {
	// if the cluster is deleted, just failed the import and let it clean up resource(the machine for lighting).
	if imp.Cluster.Status.DeletedOrCleared() && imp.Status != entity.ImportStatusFailed {
		return c.transToStatus(imp, entity.ImportStatusFailed, "cluster is deleted")
	}

	switch imp.Status {
	case entity.ImportStatusPreparing:
		return c.OnPreparing(imp)
	case entity.ImportStatusImporting:
		return c.OnImporting(imp)
	case entity.ImportStatusCanceling:
		return c.OnCanceling(imp)
	case entity.ImportStatusCanceled:
		return c.cleanUp(imp)
	case entity.ImportStatusFailed:
		return c.cleanUp(imp)
	case entity.ImportStatusCompleted:
		return c.cleanUp(imp)
	default:
		return errors.Errorf("unknown status: %v", imp.Status)
	}
}

func (c BgContext) Run() error {
	imports, err := ImportsToSync(&c.tenantAndUser.TenantID)
	if err != nil {
		return err
	}
	if len(imports) == 0 {
		return nil
	}

	c.log.Debug("imports to sync", zap.Int("number", len(imports)))

	var errs error
	for _, imp := range imports {
		imp := imp
		err := c.sync(&imp)
		if err != nil {
			errs = multierr.Append(errs, err)
		}
	}

	return errs
}

// pass nil for all tenant or with a specified tenant.
func FailImportsInRecent(tenantID *entity.ID, recentDur time.Duration) ([]entity.Import, error) {
	var imports []entity.Import
	err := db.Transaction(func(tx *gorm.DB) error {
		return entity.NewImportQuerySet(tx).PreloadCluster().
			StatusEq(entity.ImportStatusFailed).
			FinishedAtGt(time.Now().Add(-recentDur)).
			All(&imports)
	})
	if err != nil {
		return nil, errors.Annotate(err, "failed to fetch imports")
	}

	if tenantID != nil {
		var n int
		for _, imp := range imports {
			if imp.Cluster.TenantID == *tenantID {
				imports[n] = imp
				n++
			}
		}
		imports = imports[:n]
	}

	return imports, nil
}

// pass nil for all tenant or with a specified tenant.
func ImportsToSync(tenantID *entity.ID) ([]entity.Import, error) {
	var imports []entity.Import
	err := db.Transaction(func(tx *gorm.DB) error {
		return entity.NewImportQuerySet(tx.Preload("Cluster.Provider")).CleanUpEq(false).PreloadCluster().All(&imports)
	})
	if err != nil {
		return nil, errors.Annotate(err, "failed to fetch imports")
	}

	if tenantID != nil {
		var n int
		for _, imp := range imports {
			if imp.Cluster.TenantID == *tenantID {
				imports[n] = imp
				n++
			}
		}
		imports = imports[:n]
	}

	return imports, nil
}

func isTinyProfile(c *entity.Cluster) bool {
	// TODO
	// Tiny is unavailable for user to use the import feature.
	// for tiny cluster we can use more cheat machine for test internally
	// But It's OK to return false for all cluster.
	return false
}

func buildResource(c *entity.Cluster) component.Resources {
	tp := capacity.MachineType(config.Global().ImportConfig.MachineType)
	// tiny is unavailable for user to use the import feature.
	// we use c5.large just for test internal.
	if isTinyProfile(c) {
		tp = "c5.large"
	}
	return component.Resources{
		ComponentName: component.Lightning,
		//Image:
		Capacity: capacity.Capacity{
			MachineType: &tp,
		},
	}
}

func (m Model) gardenEnv() *gardener.GardEnvironment {
	gardenerDeployer := m.deployer.(gardener.GardenerDeployer)
	return gardenerDeployer.GardenerEnv
}

func (c BgContext) waitLightningStop(ctx context.Context, imp *entity.Import) error {
	cluster := imp.Cluster
	shootCli, err := c.getShootCli(cluster)
	if err != nil {
		return err
	}

	ns := runtimeutil.TiDBNamespace(cluster.ID.Uint64()).String()

	labelsMap := map[string]string{
		"job-name": name(imp),
	}

	for {
		resp, err := shootCli.CoreV1().Pods(ns).List(metav1.ListOptions{
			LabelSelector: labels.SelectorFromSet(labelsMap).String(),
		})
		if err != nil {
			c.log.Error("failed to get pod", zap.Error(err))
		}

		if len(resp.Items) == 0 {
			return nil
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(time.Second):
		}
	}
}

// cleanJob clean all resource create at shoot cluster in ensureJobCreated
func (c BgContext) cleanJob(imp *entity.Import) error {
	cluster := imp.Cluster
	shootCli, err := c.getShootCli(cluster)
	if err != nil {
		return err
	}
	ns := runtimeutil.TiDBNamespace(cluster.ID.Uint64()).String()

	background := metav1.DeletePropagationBackground
	err = shootCli.BatchV1().Jobs(ns).Delete(name(imp), &metav1.DeleteOptions{PropagationPolicy: &background})
	if err != nil && !apierrors.IsNotFound(err) {
		return errors.Annotatef(err, "failed to delete job: %s", err)
	}

	err = shootCli.CoreV1().Secrets(ns).Delete(name(imp), nil)
	if err != nil && !apierrors.IsNotFound(err) {
		return errors.Annotatef(err, "failed to delete config map: %s", err)
	}

	err = shootCli.CoreV1().Services(ns).Delete(name(imp), nil)
	if err != nil && !apierrors.IsNotFound(err) {
		return errors.Annotatef(err, "failed to delete service: %s", err)
	}

	err = shootCli.CoreV1().PersistentVolumeClaims(ns).Delete(name(imp), nil)
	if err != nil && !apierrors.IsNotFound(err) {
		return errors.Annotatef(err, "failed to delete pvc: %s", err)
	}

	return nil
}

func appendVolume(job *batchv1.Job) {
	volume := corev1.Volume{
		Name: job.Name,
		VolumeSource: corev1.VolumeSource{
			PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
				ClaimName: job.Name,
				ReadOnly:  false,
			},
		},
	}
	job.Spec.Template.Spec.Volumes = append(job.Spec.Template.Spec.Volumes, volume)

	mount := corev1.VolumeMount{
		Name:      job.Name,
		MountPath: "/etc/endpoint",
	}

	for i, c := range job.Spec.Template.Spec.Containers {
		c = job.Spec.Template.Spec.Containers[i]
		c.VolumeMounts = append(c.VolumeMounts, mount)
		job.Spec.Template.Spec.Containers[i] = c
	}
}

// ensurePVCCreated create a pvc we will use as the 'sorted-kv-dir' config item of lightning
func (c BgContext) ensurePVCCreated(imp *entity.Import) error {
	cluster := imp.Cluster
	shootCli, err := c.getShootCli(cluster)
	if err != nil {
		return err
	}
	ns := runtimeutil.TiDBNamespace(cluster.ID.Uint64()).String()

	pvc := k8s.NewPersistentVolumeClaim(name(imp), ns, config.Global().ImportConfig.StorageRequest)

	_, err = shootCli.CoreV1().PersistentVolumeClaims(ns).Create(pvc)
	if err != nil && !apierrors.IsAlreadyExists(err) {
		return errors.AddStack(err)
	}

	return nil
}

func (c BgContext) ensureJobCreated(g *gardener.GardEnvironment, imp *entity.Import) error {
	cluster := imp.Cluster
	shootCli, err := c.getShootCli(cluster)
	if err != nil {
		return err
	}
	ns := runtimeutil.TiDBNamespace(cluster.ID.Uint64()).String()

	var (
		provider *entity.Provider
		region   *entity.Region
	)
	err = db.Transaction(func(tx *gorm.DB) error {
		var err error
		provider, region, err = dbutil.GetClusterProviderAndRegion(tx, *cluster)
		return err
	})
	if err != nil {
		return err
	}

	rx, err := c.runtimes.Context(cluster.ProjectID.Uint64(), provider.Name, region.Name)
	if err != nil {
		return err
	}
	tc, err := rx.Runtime.GetTiDBCluster(ns, true)
	if err != nil {
		return err
	}

	var source string
	if strings.Contains(imp.SourceURL, "?") {
		source = fmt.Sprintf("%s&region=%s", imp.SourceURL, imp.BucketRegion)
	} else {
		source = fmt.Sprintf("%s?region=%s", imp.SourceURL, imp.BucketRegion)
	}

	var password string
	if imp.TargetPassword != "" {
		secret, err := encrypt.DecryptBase64(c.encryptor, imp.TargetPassword)
		if err != nil {
			return errors.Annotatef(err, "failed to decrypt, import id: %d", imp.ID)
		}
		password = string(secret)
	}

	cfg, err := k8s.GetLightningConfig(source, imp.TargetUsername, password, buildFilter(imp.Filter), imp.DataFormat == entity.DataFormatParquet)
	if err != nil {
		return err
	}
	image := config.Global().ImportConfig.LightningImage
	if image == "" {
		image = "pingcap/tidb-lightning:" + tc.Spec.Version
	}

	region, err = getClusterRegion(cluster)
	if err != nil {
		return err
	}
	var roleARN string
	if c.api.InfraAPIConfig.EnableNetwork {
		network, err := c.api.GetNetwork(context.TODO(), cluster.Provider.Name.String(), cluster.ProjectID.Uint64(), region.Name)
		if err != nil {
			return err
		}
		if network.Status.Import.RoleARN == "" {
			return errors.New("import role is not ready")
		}
		roleARN = network.Status.Import.RoleARN
	} else {
		shootName := runtimeutil.ShootName(cluster.ProjectID.Uint64(), cluster.Provider.Name, region.Name)
		roleARN, err = g.GetImportIamIdentity(g.ProjectName, shootName, c.tenantAndUser.TenantID.Uint64(), cluster.Provider.Name)
		if err != nil {
			return err
		}
	}

	cred := credentials(roleARN)

	job, secret := k8s.NewLightning(name(imp), ns, image, lightningNodeSelector(name(imp)), cfg, cred)

	if !isTinyProfile(cluster) {
		err = c.ensurePVCCreated(imp)
		if err != nil && !apierrors.IsAlreadyExists(err) {
			return errors.AddStack(err)
		}
		appendVolume(job)
	}

	_, err = shootCli.CoreV1().Secrets(ns).Create(secret)
	if err != nil && !apierrors.IsAlreadyExists(err) {
		return errors.AddStack(err)
	}

	_, err = shootCli.BatchV1().Jobs(ns).Create(job)
	if err != nil && !apierrors.IsAlreadyExists(err) {
		return errors.AddStack(err)
	}

	svc := k8s.NewLightningService(name(imp), ns)
	_, err = shootCli.CoreV1().Services(ns).Create(svc)
	if err != nil && !apierrors.IsAlreadyExists(err) {
		return errors.AddStack(err)
	}

	return nil
}

func (c BgContext) addAnnotationsForImport(w *gardencorev1beta1.Worker, importID uint64) {
	if w.Annotations == nil {
		w.Annotations = make(map[string]string)
	}
	w.Annotations["create-for-import-id"] = strconv.FormatUint(importID, 10)
}

func (c BgContext) isForImportWorker(w gardencorev1beta1.Worker, importID uint64) bool {
	v, ok := w.Annotations["create-for-import-id"]
	if !ok {
		return false
	}

	id, err := strconv.ParseUint(v, 10, 64)
	if err != nil {
		c.log.Error("failed to parse cluster id", zap.String("id", v), zap.Error(err))
		return false
	}

	return id == importID
}

func (c BgContext) cleanWorker(g *gardener.GardEnvironment, cluster *entity.Cluster, imp *entity.Import) error {
	region, err := getClusterRegion(cluster)
	if err != nil {
		return err
	}
	shootName := runtimeutil.ShootName(cluster.ProjectID.Uint64(), cluster.Provider.Name, region.Name)

	return retry.RetryOnConflict(retry.DefaultRetry, func() error {
		shoot, err := g.Gardencli.CoreV1beta1().Shoots(g.GetShootNamespace()).Get(shootName, metav1.GetOptions{})
		if err != nil {
			return errors.AddStack(err)
		}

		// remove the worker and update the shoot.
		var workers []gardencorev1beta1.Worker
		for _, w := range shoot.Spec.Provider.Workers {
			if c.isForImportWorker(w, identity(imp)) {
				c.log.Info("found worker to be removed", zap.Reflect("worker", w))
				continue
			}
			workers = append(workers, w)
		}

		if len(workers) != len(shoot.Spec.Provider.Workers) {
			shoot.Spec.Provider.Workers = workers
			_, err := g.Gardencli.CoreV1beta1().Shoots(g.GetShootNamespace()).Update(shoot)
			if err != nil {
				return errors.AddStack(err)
			}
			c.log.Info("clean worker success")
		}
		return nil
	})
}

func (c BgContext) getShootCli(cluster *entity.Cluster) (kubernetes.Interface, error) {
	region, err := getClusterRegion(cluster)
	if err != nil {
		return nil, err
	}

	rx, err := c.runtimes.Context(cluster.ProjectID.Uint64(), cluster.Provider.Name, region.Name)
	if err != nil {
		return nil, errors.AddStack(err)
	}

	return rx.ShootCli, nil
}

func lightningNodeSelector(name string) map[string]string {
	mp := make(map[string]string)
	mp[k8s.GardenerWorkerNodeLabelName] = name

	return mp
}

func (c BgContext) ensureWorkerCreated(g *gardener.GardEnvironment, cluster *entity.Cluster, imp *entity.Import) error {
	prefixName := runtime.GardenerWorkGroupNamePrefix(cluster.ID.Uint64(), k8s.TiDBLightningComName)

	region, err := getClusterRegion(cluster)
	if err != nil {
		return err
	}

	// TODO need some strategy to select zone?
	zone := region.Zones[0].Name
	shootName := runtimeutil.ShootName(cluster.ProjectID.Uint64(), cluster.Provider.Name, region.Name)

	var projectID = cluster.ProjectID.Uint64()
	worker, err := gardener.BuildWorkGroup(cluster.Provider.Name, name(imp), prefixName, buildResource(cluster), 1 /*vmNumber*/, zone, c.tenantAndUser.TenantID.Uint64(), projectID)
	if err != nil {
		return err
	}

	c.addAnnotationsForImport(&worker, identity(imp))

	shoot, err := g.Gardencli.CoreV1beta1().Shoots(g.GetShootNamespace()).Get(shootName, metav1.GetOptions{})
	if err != nil {
		return errors.AddStack(err)
	}

	c.log.Debug("build worker", zap.Reflect("worker", worker))
	for _, w := range shoot.Spec.Provider.Workers {
		// return if already exists the worker
		if w.Name == worker.Name {
			c.log.Debug("ignore creating worker cause already exist")
			return nil
		}
	}

	shoot.Spec.Provider.Workers = append(shoot.Spec.Provider.Workers, worker)

	_, err = g.Gardencli.CoreV1beta1().Shoots(g.GetShootNamespace()).Update(shoot)
	if err != nil {
		return errors.AddStack(err)
	}
	c.log.Info("create worker success", zap.Reflect("worker", worker))

	return nil
}

func getClusterRegion(cluster *entity.Cluster) (*entity.Region, error) {
	var regions []entity.Region
	err := db.Transaction(func(tx *gorm.DB) error {
		var err error
		regions, err = dbutil.GetClusterRegions(tx, cluster.ID)
		return err
	})
	if err != nil {
		return nil, err
	}

	if len(regions) != 1 {
		return nil, errors.Errorf("regions len should only be 1, but is %d", len(regions))
	}

	return &regions[0], nil
}

// get total files and completed percent from api directly after
// https://github.com/pingcap/tidb-lightning/issues/525 is fixed
// current we get total files from /progress/table, but the data Chunks is lazy
// initialize, so the total files will not count the pending tables.
func updateLightningInfo(imp *entity.Import, getter ProgressGetter) error {
	task, err := getter.Task()
	if err != nil {
		return err
	}

	imp.Message = task.Message
	if len(imp.Message) > 255 {
		imp.Message = imp.Message[:255]
	}

	if task.Status == taskStatusNotStarted {
		return nil
	} else if task.Status == taskStatusRunning {
		imp.Status = entity.ImportStatusImporting
	} else if task.Status == taskStatusCompleted {
		if task.Message == "" {
			imp.Status = entity.ImportStatusCompleted
		} else {
			imp.Status = entity.ImportStatusFailed
		}
	}

	var size uint64
	var totalParseSize int64
	var completed uint32
	var pendingTables uint32
	var files uint32
	runningTables := make(map[string]int64) // name to size

	for name, tableInfo := range task.Tables {
		size += uint64(tableInfo.TotalSize)
		switch tableInfo.Status {
		case taskStatusNotStarted:
			pendingTables++
		case taskStatusRunning:
			runningTables[name] = tableInfo.TotalSize
		case taskStatusCompleted:
			completed++
		}
	}
	imp.CompletedTables = completed
	imp.PendingTables = pendingTables
	imp.TotalSize = &size

	var currentTables []*entity.CurrentTable

	for name, tableInfo := range task.Tables {
		cp, err := getter.Table(name)
		if err != nil {
			log.Warn("failed to get table", zap.String("table", name), zap.Error(err))
			continue
		}

		for k, eng := range cp.Engines {
			if k == -1 { // skip index engine
				continue
			}
			files += uint32(len(eng.Chunks))
		}

		percent, parsedSize := computeCompletedPercent(cp)
		totalParseSize += parsedSize

		if _, ok := runningTables[name]; ok {
			current := &entity.CurrentTable{
				Name:             name,
				Size:             uint64(tableInfo.TotalSize),
				CompletedPercent: percent,
			}
			currentTables = append(currentTables, current)
		}
	}

	imp.CurrentTables.Tables = currentTables
	if size > 0 {
		imp.CompletedPercent = uint32(float64(totalParseSize) / float64(size) * 100)
	}

	if imp.TotalFiles == nil || *imp.TotalFiles < files {
		imp.TotalFiles = &files
	}

	return nil
}

func computeCompletedPercent(cp *TableCheckpoint) (ipercent uint32, parsedSize int64) {
	var totalSize int64

	for _, eng := range cp.Engines {
		for _, chunk := range eng.Chunks {
			totalSize += chunk.Chunk.EndOffset - chunk.Key.Offset
			parsedSize += chunk.Chunk.Offset - chunk.Key.Offset
		}
	}

	if totalSize == 0 {
		return 0, 0
	}

	// we consider the paring data frame take 80%
	// after parsing all data, we can not know the import & post processing progress
	percent := float64(parsedSize) / float64(totalSize) * 0.8

	return uint32(percent * 100), parsedSize
}

// This currently work as only one entry as the `filter` config of lightning which is a array
// see https://github.com/pingcap/tidb-lightning/blob/7622f919df0e5095daced1907b8051a75c5d6a03/tidb-lightning.toml#L146
func buildFilter(str string) (filter []string) {
	if str == "" {
		return nil
	}

	filter = append(filter, str)
	return
}

func credentials(roleARN string) string {
	// with AWS_SDK_LOAD_CONFIG setting as true
	// this file ~/.aws/credentials will make the SDK
	// get credentials from ec2 meta and than assume the specified role_arn
	format := `[default]
role_arn = %s
credential_source = Ec2InstanceMetadata
`

	return fmt.Sprintf(format, roleARN)
}