package test

import (
	"context"
	"fmt"
	"time"

	"github.com/go-logr/logr"
	"github.com/pingcap/errors"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/sets"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"

	"github.com/tidbcloud/infra-api/pkg/k8s"

	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"

	"github.com/tidbcloud/infra-api/pkg/pingcap/label"
	tidbclusterutils "github.com/tidbcloud/infra-api/pkg/pingcap/util/tidbcluster"
)

// MockTCReconciler reconciles tidbcluster object in test env
type MockTCReconciler struct {
	client.Client

	Log    logr.Logger
	Scheme *runtime.Scheme
}

func NewMockTCReconciler(mgr ctrl.Manager) *MockTCReconciler {
	return &MockTCReconciler{
		Client: mgr.GetClient(),
		Scheme: mgr.GetScheme(),
		Log:    mgr.GetLogger().WithName("MockTC"),
	}
}

func (r *MockTCReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	tc := &v1alpha1.TidbCluster{}
	err := r.Get(ctx, req.NamespacedName, tc)
	if err != nil {
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}
	return ctrl.Result{RequeueAfter: 10 * time.Second}, r.MockTidbClusterReady(ctx, r.Client, tc)
}

// SetupWithManager sets up the controller with the Manager.
func (r *MockTCReconciler) SetupWithManager(mgr ctrl.Manager) error {
	return ctrl.NewControllerManagedBy(mgr).
		For(&v1alpha1.TidbCluster{}).
		Complete(r)
}

func (r *MockTCReconciler) MockTidbClusterReady(ctx context.Context, cli client.Client, tc *v1alpha1.TidbCluster) error {
	readyCondition := tidbclusterutils.GetTidbClusterReadyCondition(tc)
	if readyCondition == nil {
		tc.Status.Conditions = append(tc.Status.Conditions, v1alpha1.TidbClusterCondition{
			Type:   v1alpha1.TidbClusterReady,
			Status: corev1.ConditionTrue,
		})
	}
	dbSvc := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Namespace: tc.Namespace,
			Name:      "db-tidb",
		},
		Spec: corev1.ServiceSpec{
			Ports: []corev1.ServicePort{{
				Name:     "test",
				Protocol: "TCP",
				Port:     10080,
			}},
		},
	}
	if tc.Spec.TiDB != nil && tc.Spec.TiDB.Service != nil && tc.Spec.TiDB.Service.Type == corev1.ServiceTypeLoadBalancer {
		dbSvc.Status = corev1.ServiceStatus{
			LoadBalancer: corev1.LoadBalancerStatus{
				Ingress: []corev1.LoadBalancerIngress{{
					IP:       "128.128.128.128",
					Hostname: "tidb.test.pingcap.com",
				}},
			},
		}
	}
	if _, err := k8s.CreateIgnoreExists(ctx, cli, dbSvc); err != nil {
		return err
	}
	if err := cli.Status().Update(ctx, dbSvc); err != nil {
		return err
	}
	zones := []string{"z1", "z2", "z3"}
	if err := r.MockComponentPodsAndNodes(ctx, cli, tc, zones, cloudv1.ComponentTiDB); err != nil {
		return err
	}
	if err := r.MockComponentPodsAndNodes(ctx, cli, tc, zones, cloudv1.ComponentPD); err != nil {
		return err
	}
	if err := r.MockComponentPodsAndNodes(ctx, cli, tc, zones, cloudv1.ComponentTiKV); err != nil {
		return err
	}
	if tc.Spec.TiDB != nil {
		dbMembers := map[string]v1alpha1.TiDBMember{}
		for i := 0; i < int(tc.Spec.TiDB.Replicas); i++ {
			dbMembers[fmt.Sprintf("db-tidb-%d", i)] = v1alpha1.TiDBMember{
				Health: true,
			}
		}
		tc.Status.TiDB = v1alpha1.TiDBStatus{
			Phase:   v1alpha1.NormalPhase,
			Members: dbMembers,
			StatefulSet: &appsv1.StatefulSetStatus{
				ReadyReplicas:   tc.Spec.TiDB.Replicas,
				CurrentReplicas: tc.Spec.TiDB.Replicas,
			},
		}
	}

	if tc.Spec.PD != nil {
		pdMembers := map[string]v1alpha1.PDMember{}
		for i := 0; i < int(tc.Spec.PD.Replicas); i++ {
			pdMembers[fmt.Sprintf("db-pd-%d", i)] = v1alpha1.PDMember{
				Health: true,
			}
		}
		tc.Status.PD = v1alpha1.PDStatus{
			Phase:   v1alpha1.NormalPhase,
			Members: pdMembers,
			StatefulSet: &appsv1.StatefulSetStatus{
				ReadyReplicas:   tc.Spec.PD.Replicas,
				CurrentReplicas: tc.Spec.PD.Replicas,
			},
		}
	}

	if tc.Spec.TiKV != nil {
		kvStore := map[string]v1alpha1.TiKVStore{}
		for i := 0; i < int(tc.Spec.TiKV.Replicas); i++ {
			kvStore[fmt.Sprintf("%d", i)] = v1alpha1.TiKVStore{
				PodName: fmt.Sprintf("db-tikv-%d", i),
				State:   v1alpha1.TiKVStateUp,
			}
		}

		tc.Status.TiKV = v1alpha1.TiKVStatus{
			Phase:  v1alpha1.NormalPhase,
			Stores: kvStore,
			StatefulSet: &appsv1.StatefulSetStatus{
				ReadyReplicas:   tc.Spec.TiKV.Replicas,
				CurrentReplicas: tc.Spec.TiKV.Replicas,
			},
		}
	}

	if tc.Spec.TiFlash != nil {
		if err := r.MockComponentPodsAndNodes(ctx, cli, tc, zones, cloudv1.ComponentTiFlash); err != nil {
			return err
		}
		flashStore := map[string]v1alpha1.TiKVStore{}
		for i := 0; i < int(tc.Spec.TiFlash.Replicas); i++ {
			flashStore[fmt.Sprintf("%d", i)] = v1alpha1.TiKVStore{
				PodName: fmt.Sprintf("db-tiflash-%d", i),
				State:   v1alpha1.TiKVStateUp,
			}
		}
		tc.Status.TiFlash = v1alpha1.TiFlashStatus{
			Phase:  v1alpha1.NormalPhase,
			Stores: flashStore,
			StatefulSet: &appsv1.StatefulSetStatus{
				ReadyReplicas:   tc.Spec.TiFlash.Replicas,
				CurrentReplicas: tc.Spec.TiFlash.Replicas,
			},
		}
	}
	return cli.Update(ctx, tc)
}

func (r *MockTCReconciler) MockComponentPodsAndNodes(ctx context.Context, cli client.Client, tc *v1alpha1.TidbCluster, zones []string, component cloudv1.ComponentType) error {
	wantedPods := sets.String{}
	var replicas int
	switch component {
	case cloudv1.ComponentPD:
		if tc.Spec.PD == nil {
			return nil
		}
		replicas = int(tc.Spec.PD.Replicas)
	case cloudv1.ComponentTiDB:
		if tc.Spec.TiDB == nil {
			return nil
		}
		replicas = int(tc.Spec.TiDB.Replicas)
	case cloudv1.ComponentTiFlash:
		if tc.Spec.TiFlash == nil {
			return nil
		}
		replicas = int(tc.Spec.TiFlash.Replicas)
	case cloudv1.ComponentTiKV:
		if tc.Spec.TiKV == nil {
			return nil
		}
		replicas = int(tc.Spec.TiKV.Replicas)
	default:
		return errors.Errorf("unknown component %s", component)
	}
	for i := 0; i < replicas; i++ {
		pod := &corev1.Pod{
			ObjectMeta: metav1.ObjectMeta{
				Name:      fmt.Sprintf("db-%s-%d", component, i),
				Namespace: tc.Namespace,
				Labels: map[string]string{
					label.ComponentLabelKey: string(component),
					label.InstanceLabelKey:  "db",
				},
			},
			Spec: corev1.PodSpec{
				Containers: []corev1.Container{{
					Name:  "test",
					Image: "test",
				}},
				NodeName: fmt.Sprintf("%s-%s-%d", tc.Namespace, component, i),
			},
		}
		if _, err := k8s.CreateIgnoreExists(ctx, cli, pod); err != nil {
			return err
		}
		wantedPods.Insert(pod.Name)
		node := &corev1.Node{
			ObjectMeta: metav1.ObjectMeta{
				Name: fmt.Sprintf("%s-%s-%d", tc.Namespace, component, i),
				Labels: map[string]string{
					"gardener.wg.prefix.name": fmt.Sprintf("%s-%s", tc.Namespace, component),
					"gardener.node.az.info":   zones[i%len(zones)],
					k8s.ZoneLabel:             zones[i%len(zones)],
				},
			},
			Status: corev1.NodeStatus{
				Conditions: []corev1.NodeCondition{
					{
						Type:   corev1.NodeReady,
						Status: corev1.ConditionTrue,
					},
				},
			},
		}
		if _, err := k8s.CreateIgnoreExists(ctx, cli, node); err != nil {
			return err
		}
		if err := cli.Status().Update(ctx, node); err != nil {
			return err
		}
	}
	// cleanup un-needed pods
	podList := &corev1.PodList{}
	err := cli.List(ctx, podList, client.InNamespace(tc.Namespace), client.MatchingLabels(
		map[string]string{
			label.ComponentLabelKey: string(component),
			label.InstanceLabelKey:  "db",
		},
	))
	if err != nil {
		return err
	}
	for _, pod := range podList.Items {
		if !wantedPods.Has(pod.Name) {
			err := cli.Delete(ctx, &pod, client.GracePeriodSeconds(0))
			if err != nil {
				return err
			}
		}
	}
	return nil
}