package branch

import (
	"context"
	"fmt"
	"os"
	"strconv"
	"time"

	"github.com/go-logr/logr"
	"github.com/tidbcloud/aws-shared-provider/pkg/manager"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/pd"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidb"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tidbworker"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/tikv"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	infrao11y "github.com/tidbcloud/observability/agent-control-server/pkg/o11y/api"
	infraclient "github.com/tidbcloud/tidb-management-service/api/serverless_infra/client"
	clientv3 "go.etcd.io/etcd/client/v3"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/client-go/tools/record"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	ctrlcluster "sigs.k8s.io/controller-runtime/pkg/cluster"
	"sigs.k8s.io/controller-runtime/pkg/controller"
	"sigs.k8s.io/controller-runtime/pkg/handler"
	"sigs.k8s.io/controller-runtime/pkg/source"
)

const (
	MaxConcurrentBranch    = 20
	CacheSyncTimeoutBranch = 5 * time.Minute
)

// Reconciler reconciles a Branch object.
type Reconciler struct {
	logr.Logger

	InfraClusterCli infraclient.Client

	Client      client.Client
	InfraClient client.Client

	InfraAPIReader client.Reader

	PdClient      pd.HttpPdClient
	TikvApi       tikv.Api
	TidbWorkerApi tidbworker.Api

	// BranchRestoreMode selects how branch keyspaces without an explicit
	// parent timestamp are restored from the parent: pitr (default) or
	// online_snapshot (the tikv worker picks the snapshot TS at task start).
	// Branches carrying a parent timestamp always restore to that point via
	// pitr regardless of the mode.
	BranchRestoreMode string

	Recorder      record.EventRecorder
	InfraRecorder record.EventRecorder

	// metric agent dependencies introduced by new o11y infra.
	O11yAgentEnabled    bool
	O11yConfig          *cloudv1.ObservabilityInfraConfig
	O11yApi             infrao11y.Api
	O11yAgentAccessArn  string
	O11yEnv             string
	O11yGlobalClusterID string

	EtcdKV                   clientv3.KV
	EtcdLease                clientv3.Lease
	S3BucketName             string
	BaseProxy                string
	KubeClusterName          string
	KeyspacePoolNs           string
	EncryptKeyspacePoolNs    string
	EtcdTokenPrefixMap       map[string]string
	TidbVersionMap           map[string]string
	TidbPools                []string
	ProviderType             string
	ServerlessLightningImage string
	TidbManagers             map[string]tidbmanager.TidbManager
	TidbController           tidb.Controller

	SkipClusterDeleteEtcdAndPd bool
}

// reconcileContext is the local context of a single branch reconciliation, it is NOT concurrency safe.
type reconcileContext struct {
	*Reconciler
	logr.Logger

	ctx context.Context
	// the branch will be continuously updated during reconciliation.
	branch *cloudv1.Branch

	keyspaceNs    string
	keyspaceToken string
	keyspaceID    uint32

	specTidbPool, etcdTokenPrefix, tidbVersion string
}

func (r *Reconciler) newReconcileContext(ctx context.Context, req ctrl.Request) (*reconcileContext, error) {
	log := ctrl.LoggerFrom(ctx)

	branch := &cloudv1.Branch{}
	if err := r.InfraClient.Get(ctx, req.NamespacedName, branch); err != nil {
		if apierrors.IsNotFound(err) {
			return nil, nil
		}
		return nil, fmt.Errorf("can't get branch: %w", err)
	}

	if kubeCluster := labels.GetKubeClusterFromLabels(branch.Labels); kubeCluster != r.KubeClusterName {
		log.Info("unexpected branch scheduling", "kubecluster", kubeCluster)
		return nil, nil
	}

	clusterStatus, err := branch.AwsFreeTierClusterStatusFromBranch()
	if err != nil {
		return nil, fmt.Errorf("AwsFreeTierClusterStatusFromBranch err: %w", err)
	}

	tidbPool := r.getTidbPool(branch)

	// we must get the latest branch from infra in some conditions.
	if clusterStatus.KeyspaceToken == "" ||
		tidbPool != clusterStatus.TidbPool || branch.Annotations[types.AnnoKeyTidbPoolChanging] != "" {
		latestBranch := &cloudv1.Branch{}
		if err := r.InfraAPIReader.Get(ctx, req.NamespacedName, latestBranch); err != nil {
			return nil, fmt.Errorf("can't get branch from APIReader: %w", err)
		}
		if latestBranch.ResourceVersion != branch.ResourceVersion {
			return nil, fmt.Errorf("branch resourceVersion(%v) is not the latest(%v), reconcile again",
				branch.ResourceVersion, latestBranch.ResourceVersion)
		}
	}

	log = log.WithValues("rrv", branch.ResourceVersion)
	if clusterStatus.KeyspaceToken != "" {
		log = log.WithValues("keyspaceID", clusterStatus.KeyspaceId)
	}

	keyspaceNs := r.KeyspacePoolNs
	if branch.Labels[labels.LabelEncryptKeyspace] == "true" {
		keyspaceNs = r.EncryptKeyspacePoolNs
	}

	// Branches default to a single TiDB replica, can be overridden via forProvider.
	tidbMinReplicas := 1
	tidbMaxReplicas := 1
	specForProvider, err := types.InfraServerlessSpecForProvider(branch)
	if err != nil {
		return nil, fmt.Errorf("can't get spec.forProvider: %w", err)
	}
	if specForProvider.TidbMaxReplicas > 0 {
		tidbMaxReplicas = int(specForProvider.TidbMaxReplicas)
		log.Info("use infra spec.forProvider", "tidbMaxReplicas", specForProvider.TidbMaxReplicas)
	}
	if specForProvider.TidbMinReplicas > 0 {
		tidbMinReplicas = int(specForProvider.TidbMinReplicas)
		log.Info("use infra spec.forProvider", "tidbMinReplicas", specForProvider.TidbMinReplicas)
	}
	if tidbMinReplicas < 1 {
		log.Info("tidbMinReplicas < 1, reset tidbMinReplicas to defaultMinReplicas", "tidbMinReplicas", tidbMinReplicas)
		tidbMinReplicas = 1
	}
	if tidbMinReplicas > tidbMaxReplicas {
		log.Info("tidbMinReplicas > tidbMaxReplicas, reset tidbMinReplicas to defaultMinReplicas",
			"tidbMinReplicas", tidbMinReplicas, "tidbMaxReplicas", tidbMaxReplicas)
		tidbMinReplicas = 1
	}
	if branch.Spec.Cluster.TiDB.Replicas > tidbMaxReplicas {
		log.Info("tidb replicas is greater than tidbMaxReplicas, force reset tidb replicas", "tidbPool", tidbPool,
			"tidb replicas", branch.Spec.Cluster.TiDB.Replicas, "tidbMaxReplicas", tidbMaxReplicas,
		)
		branch.Spec.Cluster.TiDB.Replicas = tidbMaxReplicas
		if err := r.InfraClient.Update(ctx, branch); err != nil {
			return nil, fmt.Errorf("update branch for tidb replicas err: %w", err)
		}
	} else if branch.Spec.Cluster.TiDB.Replicas < tidbMinReplicas {
		log.Info("tidb replicas is less than tidbMinReplicas, force reset tidb replicas", "tidbPool", tidbPool,
			"tidb replicas", branch.Spec.Cluster.TiDB.Replicas, "tidbMinReplicas", tidbMinReplicas,
		)
		branch.Spec.Cluster.TiDB.Replicas = tidbMinReplicas
		if err := r.InfraClient.Update(ctx, branch); err != nil {
			return nil, fmt.Errorf("update branch for tidb replicas err: %w", err)
		}
	}

	return &reconcileContext{
		Reconciler: r,
		Logger:     log,

		ctx:             ctx,
		branch:          branch,
		keyspaceNs:      keyspaceNs,
		keyspaceToken:   clusterStatus.KeyspaceToken,
		keyspaceID:      clusterStatus.KeyspaceId,
		specTidbPool:    tidbPool,
		etcdTokenPrefix: r.EtcdTokenPrefixMap[tidbPool],
		tidbVersion:     r.TidbVersionMap[tidbPool],
	}, nil
}

func (r *Reconciler) getTidbPool(branch *cloudv1.Branch) string {
	tidbPool := r.TidbPools[0] // The first element is default TidbPool
	for _, v := range r.TidbPools {
		if v == branch.Annotations[types.AnnoTidbPoolSelector] {
			tidbPool = branch.Annotations[types.AnnoTidbPoolSelector]
			break
		}
	}

	return tidbPool
}

func (r *reconcileContext) needResetTidbReplicas() bool {
	if r.branch.Status.DB == nil || r.branch.Status.DB.TiDB == nil {
		return true
	}
	clusterStatus, err := r.branch.AwsFreeTierClusterStatusFromBranch()
	if err != nil {
		r.Error(err, "call AwsFreeTierClusterStatusFromBranch err")
		return true
	}
	if r.specTidbPool != clusterStatus.TidbPool {
		r.Info("tidb pool have changed", "old", clusterStatus.TidbPool, "now", r.specTidbPool)
		return true
	}
	if int32(r.branch.Spec.Cluster.TiDB.Replicas) != r.branch.Status.DB.TiDB.UpdatedReplicas {
		r.Info("tidb replicas have changed",
			"old", r.branch.Status.DB.TiDB.UpdatedReplicas, "now", r.branch.Spec.Cluster.TiDB.Replicas)
		return true
	}

	return false
}

// Reconcile is part of the main kubernetes reconciliation loop which aims to
// move the current state of the cluster closer to the desired state.
func (r *Reconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	log := ctrl.LoggerFrom(ctx)
	log.Info("start reconcile branch")
	start := time.Now()
	defer func() {
		log.Info("finish reconcile branch", "time", time.Since(start))
	}()

	rx, err := r.newReconcileContext(ctx, req)
	if err != nil || rx == nil {
		return utils.ReconcileErr(log, "Branch", req.Name, err)
	}

	if err = rx.ensureFinalizer(); err != nil {
		return utils.ReconcileErr(log, "Branch", rx.branch.Name, err)
	}

	_, err = r.shortcut(rx)
	if err != nil {
		return utils.ReconcileErr(log, "Branch", rx.branch.Name, err)
	}

	branchPhase := rx.branch.Status.Phase
	switch branchPhase {
	case cloudv1.BranchPhaseCreating, cloudv1.BranchPhasePending:
		err = rx.createBranch()
	case cloudv1.BranchPhaseRunning:
		err = rx.reconcileBranch()
	case cloudv1.BranchPhaseRestoring:
		err = rx.restoreBranch()
	case cloudv1.BranchPhasePendingDelete:
		err = rx.pendingDeleteBranch()
	case cloudv1.BranchPhaseDeleting:
		err = rx.deleteBranch()
	case cloudv1.BranchPhaseUnavailable:
		// when branch is unavailable, skip reconciliation and observe until it recovers for safety
		err = rx.observeBranch()
	default:
		// unknown phase, try to recover the real phase via observation
		err = rx.observeBranch()
	}

	if branchPhase == cloudv1.BranchPhaseDeleting && err == nil {
		err = rx.tryRemoveFinalizer()
	}

	return utils.ReconcileErr(log, "Branch", rx.branch.Name, err)
}

func (r *Reconciler) shortcut(rx *reconcileContext) (ctrl.Result, error) { //nolint:unparam // by design
	var phase cloudv1.BranchPhase
	var reason string
	var needUpdatePhase bool

	if rx.branch.DeletionTimestamp != nil {
		if rx.branch.Status.Phase != cloudv1.BranchPhaseDeleting &&
			rx.branch.Status.Phase != cloudv1.BranchPhasePendingDelete {
			// shortcut to PendingDelete phase whenever cluster is get deleted
			phase = cloudv1.BranchPhasePendingDelete
			reason = "branch is pending to be deleted"
			needUpdatePhase = true
		}
	} else {
		if rx.branch.Status.Phase == "" {
			// shortcut to Pending phase when branch is newly created
			phase = cloudv1.BranchPhasePending
			reason = "cluster is waiting to be reconciled"
			needUpdatePhase = true
		}
	}

	if needUpdatePhase {
		if err := rx.updateProgress(phase, reason); err != nil {
			return utils.ReconcileErr(rx.Logger, "Branch", rx.branch.Name, err)
		}
	}
	return ctrl.Result{}, nil
}

// SetupWithManager sets up the controller with the Manager.
func (r *Reconciler) SetupWithManager(mgr manager.Manager) error {
	return r.setupWithManager(mgr, mgr)
}

func (r *Reconciler) setupWithManager(mgr manager.Manager, branchCRCluster ctrlcluster.Cluster) error {
	r.Logger.Info("start to set up branch controller with manager")
	r.Client = mgr.GetClient()
	r.InfraClient = branchCRCluster.GetClient()
	r.InfraAPIReader = branchCRCluster.GetAPIReader()
	r.Recorder = mgr.GetEventRecorderFor("Branch")
	r.InfraRecorder = branchCRCluster.GetEventRecorderFor("Branch")

	maxConcurrent := MaxConcurrentBranch
	if v := os.Getenv("MAX_CONCURRENT_CLUSTERS"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 {
			maxConcurrent = n
		}
	}

	controllerOptions := controller.Options{
		MaxConcurrentReconciles: maxConcurrent,
		Reconciler:              r,
		CacheSyncTimeout:        CacheSyncTimeoutBranch,
	}

	c, err := controller.New("Branch", mgr, controllerOptions)
	if err != nil {
		return fmt.Errorf("can't new a controller: %w", err)
	}
	r.Logger = c.GetLogger()

	if err := branchCRCluster.GetFieldIndexer().IndexField(context.TODO(), &cloudv1.Branch{},
		"status.fromProvider.keyspaceToken", func(obj client.Object) []string {
			branch := obj.(*cloudv1.Branch)
			clusterStatus, err := branch.AwsFreeTierClusterStatusFromBranch()
			if err != nil {
				r.Error(err, "get keyspace token", "branch", branch.Name)
				return nil
			}
			if clusterStatus.KeyspaceToken != "" {
				return []string{clusterStatus.KeyspaceToken}
			}
			return nil
		}); err != nil {
		return fmt.Errorf("unable to add index for branch token: %w", err)
	}
	if err := branchCRCluster.GetFieldIndexer().IndexField(context.TODO(), &cloudv1.Branch{},
		"status.fromProvider.zoneRedundantEnabled", func(obj client.Object) []string {
			clusterStatus, err := obj.(*cloudv1.Branch).AwsFreeTierClusterStatusFromBranch()
			if err != nil {
				r.Error(err, "get clusterStatus", "branch", obj.GetName())
				return nil
			}
			if clusterStatus.ZoneRedundantEnabled {
				return []string{"zoneRedundantEnabled"}
			}
			return nil
		}); err != nil {
		return fmt.Errorf("unable to add index for branch zoneRedundantEnabled: %w", err)
	}

	return r.watches(c, branchCRCluster)
}

func (r *Reconciler) watches(c controller.Controller, branchCRCluster ctrlcluster.Cluster) error {
	if err := c.Watch(source.Kind(branchCRCluster.GetCache(), &cloudv1.Branch{},
		&handler.TypedEnqueueRequestForObject[*cloudv1.Branch]{},
	)); err != nil {
		return fmt.Errorf("can't watch cluster: %w", err)
	}

	return nil
}