package branch

import (
	"time"

	"github.com/pingcap/errors"
	corev1 "k8s.io/api/core/v1"

	"github.com/tidbcloud/aws-shared-provider/pkg/metrics"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/pkg/consts"
)

const (
	defaultKillAllTimeout = 5 * time.Second
)

func (r *reconcileContext) restoreBranch() error {
	stage := r.branch.Annotations[types.AnnoKeyKeyspaceRestore]
	switch stage {
	case types.AnnoKeyspaceRestoreStart:
		return r.restoreBranchStart()
	case types.AnnoKeyspaceRestartTiDB:
		return r.restoreKeyspace()
	case types.AnnoKeyspaceRestoreFinish:
		return r.observeRestore()
	}
	return nil
}

func (r *reconcileContext) restoreBranchStart() error {
	r.Info("start to restore branch")
	// kill all TiDB connections
	for pool, tManager := range r.TidbManagers {
		if err := tManager.KillAll(r.ctx, r.keyspaceToken, defaultKillAllTimeout); err != nil {
			return errors.Annotatef(err, "failed to kill all tidbs for keyspaceToken in pool %s", pool)
		}
	}
	r.Info("reconcileResetStatus kill all tidb succeed")

	// stop all tidb workers, after killall connections and before set bootstrap flag,
	// to avoid any tidb instance starts and trigger user amend operation.
	if err := r.TidbWorkerApi.Clean(r.ctx, r.keyspaceToken); err != nil {
		return errors.Annotate(err, "failed to clean tidb worker")
	}

	r.branch.Annotations[types.AnnoKeyKeyspaceRestore] = types.AnnoKeyspaceRestartTiDB
	r.branch.Annotations[types.AnnoKeyRestoreStartTimestamp] = time.Now().UTC().Format(time.RFC3339)
	// mark branch does not have restore-id, so that `needRestoreKeyspace` will return true
	delete(r.branch.Annotations, types.AnnoKeyBranchRestoreId)
	delete(r.branch.Annotations, types.AnnoKeyBranchErrorAndRecreate)
	if err := r.InfraClient.Update(r.ctx, r.branch); err != nil {
		return errors.Annotate(err, "failed to update infra branch")
	}
	return nil
}

func (r *reconcileContext) restoreKeyspace() error {
	if r.needRestoreKeyspace() {
		if err := r.updateProgress(cloudv1.BranchPhaseRestoring, "starting restore keyspace"); err != nil {
			return err
		}
		if err := r.createRestoreKeyspace(); err != nil {
			return err
		}
	}
	if ok, err := r.checkRestoreKeyspace(); err != nil {
		return errors.Annotate(err, "failed to check restore keyspace")
	} else if !ok {
		if r.branch.Status.Reason == "still restoring keyspace" {
			return &utils.RequeueAfterError{
				Cause:        nil,
				RequeueAfter: time.Second * 5, //nolint:mnd // mask
			}
		}
		if err := r.updateProgress(cloudv1.BranchPhaseRestoring, "still restoring keyspace"); err != nil {
			return err
		}
		return nil
	}

	if err := r.setKeyspaceConfigIsBranchBootstrapped(); err != nil {
		return errors.Annotate(err, "failed to set IsBranchBootstrapped")
	}

	if err := r.ensureResourceGroup(); err != nil {
		return err
	}

	if err := r.ensureZoneRedundant(); err != nil {
		return err
	}

	r.branch.Annotations[types.AnnoKeyKeyspaceRestore] = types.AnnoKeyspaceRestoreFinish
	if err := r.InfraClient.Update(r.ctx, r.branch); err != nil {
		return errors.Annotate(err, "failed to update AnnoKeyKeyspaceRestore")
	}
	return nil
}

// set branch.phase to Running and ensure all other necessary status together
func (r *reconcileContext) observeRestore() error {
	if err := r.observeTidbClusterStatus(); err != nil {
		return err
	}

	r.InfraRecorder.Event(r.branch, corev1.EventTypeNormal, string(r.branch.Status.Phase), r.branch.Status.Reason)

	start := r.branch.Annotations[types.AnnoKeyRestoreStartTimestamp]
	startTime, err := time.Parse(time.RFC3339, start)
	if err != nil {
		return errors.Annotate(err, "failed to parse restore start time")
	}
	metrics.BranchRestoreHistogram.WithLabelValues(
		r.branch.Name,
		r.branch.Labels[consts.ProjectTagKey],
		r.branch.Labels[consts.TenantTagKey],
	).Observe(time.Since(startTime).Seconds())

	return nil
}