package cluster

import (
	"context"
	"math/rand"
	"time"

	"github.com/pingcap/errors"
	serverlessv1alpha1 "github.com/tidbcloud/aws-shared-provider/apis/serverless/v1alpha1"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

func (r *reconcileContext) assignKeyspace() error {
	keyspace, err := r.getKeyspaceIfBindCluster()
	if err != nil {
		return errors.Annotate(err, "getKeyspaceIfBindCluster error")
	}

	if r.keyspaceToken != "" {
		if keyspace == nil {
			return errors.New("keyspaceToken already set in cluster, but Keyspace not found")
		}
		if keyspace.Status.Phase != serverlessv1alpha1.BindPhase {
			err = r.bindKeyspace(keyspace, serverlessv1alpha1.BindPhase)
		}
		return err
	}

	if keyspace == nil {
		keyspaceInPool := &serverlessv1alpha1.KeyspaceList{}
		if err := r.Client.List(r.ctx, keyspaceInPool,
			client.InNamespace(r.keyspaceNs), client.MatchingFields{"status.inPool": "true"}); err != nil {
			return err
		}
		availableKeyspaces := make([]*serverlessv1alpha1.Keyspace, 0, len(keyspaceInPool.Items))
		for i := range keyspaceInPool.Items {
			if keyspaceInPool.Items[i].Status.Phase == serverlessv1alpha1.AvailablePhase &&
				keyspaceInPool.Items[i].DeletionTimestamp == nil {
				availableKeyspaces = append(availableKeyspaces, &keyspaceInPool.Items[i])
			}
		}
		if len(availableKeyspaces) == 0 {
			return errors.New("not enough keyspace for assign")
		}

		keyspace = availableKeyspaces[rand.Intn(len(availableKeyspaces))]

		if err := r.bindKeyspace(keyspace, serverlessv1alpha1.BindingPhase); err != nil {
			return errors.Annotate(err, "failed to Binding keyspace")
		}
	}

	if err := r.ensureStatus(func(status *infra.ClusterStatus) error {
		status.KeyspaceToken = keyspace.Spec.KeyspaceName
		status.KeyspaceId = keyspace.Status.KeyspaceId
		return nil
	}); err != nil {
		r.waitEnsureKeyspaceBinding(keyspace)
		return errors.Annotate(err, "failed to update keyspaceToken")
	}
	r.keyspaceToken = keyspace.Spec.KeyspaceName
	r.keyspaceID = keyspace.Status.KeyspaceId

	if err := r.bindKeyspace(keyspace, serverlessv1alpha1.BindPhase); err != nil {
		return errors.Annotate(err, "failed to Bind keyspace")
	}

	r.Logger = r.Logger.WithValues("keyspaceToken", r.keyspaceToken)
	return nil
}

func (r *reconcileContext) getKeyspaceIfBindCluster() (*serverlessv1alpha1.Keyspace, error) {
	bindKeyspaceList := &serverlessv1alpha1.KeyspaceList{}
	err := r.Client.List(r.ctx, bindKeyspaceList,
		client.InNamespace(r.keyspaceNs), client.MatchingFields{"status.bindCluster": r.cluster.Name})
	if err != nil {
		return nil, err
	}
	if len(bindKeyspaceList.Items) == 0 {
		return nil, nil
	}
	if len(bindKeyspaceList.Items) > 1 {
		return nil, errors.Errorf("cluster bind %v Keyspaces!", len(bindKeyspaceList.Items))
	}
	return &bindKeyspaceList.Items[0], nil
}

func (r *reconcileContext) bindKeyspace(
	keyspace *serverlessv1alpha1.Keyspace,
	bindPhase serverlessv1alpha1.PhaseStatus,
) error {
	keyspace.Status.Phase = bindPhase
	keyspace.Status.InPool = false
	keyspace.Status.BindCluster = r.cluster.Name
	keyspace.Status.ClusterNs = r.cluster.Namespace
	err := r.Client.Status().Update(r.ctx, keyspace)
	if err == nil && bindPhase == serverlessv1alpha1.BindPhase {
		r.Recorder.Event(keyspace, corev1.EventTypeNormal, "bind keyspace", "bind keyspace success")
	}
	return err
}

func (r *reconcileContext) waitEnsureKeyspaceBinding(keyspace *serverlessv1alpha1.Keyspace) {
	_ = wait.PollUntilContextCancel(r.ctx, time.Second, true, func(ctx context.Context) (done bool, err error) {
		newKeyspace := &serverlessv1alpha1.Keyspace{}
		if err = r.Client.Get(ctx, client.ObjectKeyFromObject(keyspace), newKeyspace); err != nil {
			return
		}
		return newKeyspace.Status.Phase == serverlessv1alpha1.BindingPhase, nil
	})
}