// Copyright 2026 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//go:build e2e

package waiter

import (
	"context"
	"fmt"
	"time"

	"github.com/onsi/ginkgo/v2"
	tioperatorv1alpha1 "github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	k8swait "k8s.io/apimachinery/pkg/util/wait"
	metricsclient "k8s.io/metrics/pkg/client/clientset/versioned"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/tidbcloud/cluster-service-ng/internal/capacity-planner/test/e2e/data"
)

// ClusterObserved waits for TiDB Operator to observe the current generation.
func ClusterObserved(
	ctx context.Context,
	c client.Client,
	cluster *tioperatorv1alpha1.Cluster,
) error {
	last := &tioperatorv1alpha1.Cluster{}
	err := k8swait.PollUntilContextTimeout(ctx, PollInterval, ResourceTimeout, true, func(ctx context.Context) (bool, error) {
		if err := c.Get(ctx, client.ObjectKeyFromObject(cluster), last); err != nil {
			return false, err
		}
		return last.Status.ObservedGeneration == last.Generation, nil
	})
	if err != nil {
		return fmt.Errorf("wait for TiDB Operator to reconcile Cluster %s/%s: %w; last object: %+v",
			cluster.Namespace, cluster.Name, err, last)
	}
	logf("TiDB Operator reconciled Cluster %s/%s at generation %d",
		cluster.Namespace, cluster.Name, last.Generation)
	return nil
}

type groupState struct {
	observedGeneration int64
	replicas           int32
	readyReplicas      int32
	updatedReplicas    int32
	currentRevision    string
	updateRevision     string
	conditions         []metav1.Condition
}

// PDGroupReady waits until a real PDGroup is fully rolled out.
func PDGroupReady(
	ctx context.Context,
	c client.Client,
	group *tioperatorv1alpha1.PDGroup,
	expectedReplicas int32,
) error {
	return typedGroupReady(ctx, c, group, &tioperatorv1alpha1.PDGroup{}, "PDGroup", expectedReplicas,
		func(current *tioperatorv1alpha1.PDGroup) groupState {
			return groupState{
				observedGeneration: current.Status.ObservedGeneration,
				replicas:           current.Status.Replicas,
				readyReplicas:      current.Status.ReadyReplicas,
				updatedReplicas:    current.Status.UpdatedReplicas,
				currentRevision:    current.Status.CurrentRevision,
				updateRevision:     current.Status.UpdateRevision,
				conditions:         current.Status.Conditions,
			}
		})
}

// TiKVGroupReady waits until a real TiKVGroup is fully rolled out.
func TiKVGroupReady(
	ctx context.Context,
	c client.Client,
	group *tioperatorv1alpha1.TiKVGroup,
	expectedReplicas int32,
) error {
	return typedGroupReady(ctx, c, group, &tioperatorv1alpha1.TiKVGroup{}, "TiKVGroup", expectedReplicas,
		func(current *tioperatorv1alpha1.TiKVGroup) groupState {
			return groupState{
				observedGeneration: current.Status.ObservedGeneration,
				replicas:           current.Status.Replicas,
				readyReplicas:      current.Status.ReadyReplicas,
				updatedReplicas:    current.Status.UpdatedReplicas,
				currentRevision:    current.Status.CurrentRevision,
				updateRevision:     current.Status.UpdateRevision,
				conditions:         current.Status.Conditions,
			}
		})
}

// TiDBGroupReady waits until a real TiDBGroup is fully rolled out.
func TiDBGroupReady(
	ctx context.Context,
	c client.Client,
	group *tioperatorv1alpha1.TiDBGroup,
	expectedReplicas int32,
) error {
	return typedGroupReady(ctx, c, group, &tioperatorv1alpha1.TiDBGroup{}, data.TiDBGroupKind, expectedReplicas,
		func(current *tioperatorv1alpha1.TiDBGroup) groupState {
			return groupState{
				observedGeneration: current.Status.ObservedGeneration,
				replicas:           current.Status.Replicas,
				readyReplicas:      current.Status.ReadyReplicas,
				updatedReplicas:    current.Status.UpdatedReplicas,
				currentRevision:    current.Status.CurrentRevision,
				updateRevision:     current.Status.UpdateRevision,
				conditions:         current.Status.Conditions,
			}
		})
}

func typedGroupReady[T client.Object](
	ctx context.Context,
	c client.Client,
	group, current T,
	kind string,
	expectedReplicas int32,
	stateFor func(T) groupState,
) error {
	err := k8swait.PollUntilContextTimeout(ctx, time.Second, ResourceTimeout, true, func(ctx context.Context) (bool, error) {
		if err := c.Get(ctx, client.ObjectKeyFromObject(group), current); err != nil {
			return false, err
		}
		state := stateFor(current)
		return state.observedGeneration == current.GetGeneration() &&
			state.replicas == expectedReplicas &&
			state.readyReplicas == expectedReplicas &&
			state.updatedReplicas == expectedReplicas &&
			state.currentRevision != "" &&
			state.currentRevision == state.updateRevision &&
			meta.IsStatusConditionTrue(state.conditions, tioperatorv1alpha1.CondSynced) &&
			meta.IsStatusConditionTrue(state.conditions, tioperatorv1alpha1.CondReady), nil
	})
	if err != nil {
		return fmt.Errorf("wait for real %s %s/%s to become Ready with %d replica(s): %w; last object: %+v",
			kind, group.GetNamespace(), group.GetName(), expectedReplicas, err, current)
	}
	logf("real %s %s/%s is Ready with %d replica(s)",
		kind, group.GetNamespace(), group.GetName(), expectedReplicas)
	return nil
}

// ClusterAvailable waits until TiDB Operator reports the Cluster Available.
func ClusterAvailable(
	ctx context.Context,
	c client.Client,
	cluster *tioperatorv1alpha1.Cluster,
) error {
	last := &tioperatorv1alpha1.Cluster{}
	err := k8swait.PollUntilContextTimeout(ctx, time.Second, ResourceTimeout, true, func(ctx context.Context) (bool, error) {
		if err := c.Get(ctx, client.ObjectKeyFromObject(cluster), last); err != nil {
			return false, err
		}
		return last.Status.PD != "" &&
			meta.IsStatusConditionTrue(last.Status.Conditions, tioperatorv1alpha1.ClusterCondAvailable), nil
	})
	if err != nil {
		return fmt.Errorf("wait for real TiDB Cluster %s/%s to become Available: %w; last object: %+v",
			cluster.Namespace, cluster.Name, err, last)
	}
	logf("real TiDB Cluster %s/%s is Available", cluster.Namespace, cluster.Name)
	return nil
}

// ComponentPodReady waits for a Ready Pod belonging to a TiDB component.
func ComponentPodReady(
	ctx context.Context,
	c client.Client,
	namespace, group, component string,
) (*corev1.Pod, error) {
	var last corev1.PodList
	err := k8swait.PollUntilContextTimeout(ctx, time.Second, ResourceTimeout, true, func(ctx context.Context) (bool, error) {
		last = corev1.PodList{}
		if err := c.List(ctx, &last,
			client.InNamespace(namespace),
			client.MatchingLabels{
				tioperatorv1alpha1.LabelKeyGroup:     group,
				tioperatorv1alpha1.LabelKeyComponent: component,
			},
		); err != nil {
			return false, err
		}
		for i := range last.Items {
			if IsPodReady(&last.Items[i]) {
				return true, nil
			}
		}
		return false, nil
	})
	if err != nil {
		return nil, fmt.Errorf("wait for real %s Pod in group %s/%s: %w; last Pods: %v",
			component, namespace, group, err, last.Items)
	}
	for i := range last.Items {
		if IsPodReady(&last.Items[i]) {
			return last.Items[i].DeepCopy(), nil
		}
	}
	return nil, fmt.Errorf("real %s Pod in group %s/%s disappeared after becoming Ready", component, namespace, group)
}

// TiDBGroupDesiredState waits until the autoscaler has applied the exact CPU
// tier and replica count to the target spec. It does not wait for rollout.
func TiDBGroupDesiredState(
	ctx context.Context,
	c client.Client,
	target *tioperatorv1alpha1.TiDBGroup,
	expectedCPU string,
	expectedReplicas int32,
) error {
	wantCPU := resource.MustParse(expectedCPU)
	last := &tioperatorv1alpha1.TiDBGroup{}
	err := k8swait.PollUntilContextTimeout(ctx, PollInterval, ResourceTimeout, true, func(ctx context.Context) (bool, error) {
		if err := c.Get(ctx, client.ObjectKeyFromObject(target), last); err != nil {
			return false, err
		}
		return tiDBGroupMatches(last, &wantCPU, expectedReplicas, false), nil
	})
	if err != nil {
		return fmt.Errorf("wait for TiDBGroup %s/%s desired CPU %s with %d replica(s): %w; last object: %+v",
			target.Namespace, target.Name, expectedCPU, expectedReplicas, err, last)
	}
	logf("TiDBGroup %s/%s desired state is CPU %s with %d replica(s)",
		target.Namespace, target.Name, expectedCPU, expectedReplicas)
	return nil
}

// TiDBGroupConverged waits until the target has the exact requested CPU and
// replica count and TiDB Operator reports that generation fully rolled out.
func TiDBGroupConverged(
	ctx context.Context,
	c client.Client,
	target *tioperatorv1alpha1.TiDBGroup,
	expectedCPU string,
	expectedReplicas int32,
) error {
	wantCPU := resource.MustParse(expectedCPU)
	last := &tioperatorv1alpha1.TiDBGroup{}
	err := k8swait.PollUntilContextTimeout(ctx, PollInterval, ResourceTimeout, true, func(ctx context.Context) (bool, error) {
		if err := c.Get(ctx, client.ObjectKeyFromObject(target), last); err != nil {
			return false, err
		}
		return tiDBGroupMatches(last, &wantCPU, expectedReplicas, true), nil
	})
	if err != nil {
		return fmt.Errorf("wait for TiDBGroup %s/%s to converge to CPU %s with %d replica(s): %w; last object: %+v",
			target.Namespace, target.Name, expectedCPU, expectedReplicas, err, last)
	}
	logf("TiDBGroup %s/%s converged to CPU %s with %d Ready replica(s)",
		target.Namespace, target.Name, expectedCPU, expectedReplicas)
	return nil
}

// TiDBGroupStable verifies continuously that a converged target stays at the
// exact requested state across multiple local autoscaler reconcile periods.
func TiDBGroupStable(
	ctx context.Context,
	c client.Client,
	target *tioperatorv1alpha1.TiDBGroup,
	expectedCPU string,
	expectedReplicas int32,
	duration time.Duration,
) error {
	wantCPU := resource.MustParse(expectedCPU)
	deadline := time.NewTimer(duration)
	defer deadline.Stop()
	ticker := time.NewTicker(PollInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return fmt.Errorf("verify TiDBGroup convergence: %w", ctx.Err())
		case <-deadline.C:
			return nil
		case <-ticker.C:
			current := &tioperatorv1alpha1.TiDBGroup{}
			if err := c.Get(ctx, client.ObjectKeyFromObject(target), current); err != nil {
				return fmt.Errorf("get converged TiDBGroup: %w", err)
			}
			if !tiDBGroupMatches(current, &wantCPU, expectedReplicas, true) {
				return fmt.Errorf("TiDBGroup left converged CPU %s / replicas %d state: %+v",
					expectedCPU, expectedReplicas, current)
			}
		}
	}
}

func tiDBGroupMatches(
	current *tioperatorv1alpha1.TiDBGroup,
	wantCPU *resource.Quantity,
	expectedReplicas int32,
	requireReady bool,
) bool {
	cpu := current.Spec.Template.Spec.Resources.CPU
	if ptr.Deref(current.Spec.Replicas, 0) != expectedReplicas ||
		cpu == nil || cpu.Cmp(*wantCPU) != 0 ||
		current.Labels[data.TierLabelKey] != wantCPU.String() ||
		current.Labels[data.PreserveLabelKey] != "true" {
		return false
	}
	if !requireReady {
		return true
	}
	return current.Status.ObservedGeneration == current.Generation &&
		current.Status.Replicas == expectedReplicas &&
		current.Status.ReadyReplicas == expectedReplicas &&
		current.Status.UpdatedReplicas == expectedReplicas &&
		current.Status.CurrentRevision != "" &&
		current.Status.CurrentRevision == current.Status.UpdateRevision &&
		meta.IsStatusConditionTrue(current.Status.Conditions, tioperatorv1alpha1.CondSynced) &&
		meta.IsStatusConditionTrue(current.Status.Conditions, tioperatorv1alpha1.CondReady)
}

// TiDBGroupCPU waits for fresh metrics for every Ready, non-terminating TiDB
// Pod and verifies their aggregate CPU is within the requested closed range.
// A negative maximum disables the upper bound.
func TiDBGroupCPU(
	ctx context.Context,
	c client.Client,
	metrics metricsclient.Interface,
	namespace, group string,
	expectedReplicas int32,
	minimumMilliCPU, maximumMilliCPU int64,
	notBefore time.Time,
) (int64, time.Time, error) {
	var (
		observedMilliCPU int64
		oldestSample     time.Time
		lastPods         corev1.PodList
	)
	err := k8swait.PollUntilContextTimeout(ctx, time.Second, ResourceTimeout, true, func(ctx context.Context) (bool, error) {
		lastPods = corev1.PodList{}
		if err := c.List(ctx, &lastPods,
			client.InNamespace(namespace),
			client.MatchingLabels{
				tioperatorv1alpha1.LabelKeyGroup:     group,
				tioperatorv1alpha1.LabelKeyComponent: tioperatorv1alpha1.LabelValComponentTiDB,
			},
		); err != nil {
			return false, err
		}

		live := make([]*corev1.Pod, 0, len(lastPods.Items))
		for i := range lastPods.Items {
			pod := &lastPods.Items[i]
			if pod.DeletionTimestamp.IsZero() && IsPodReady(pod) {
				live = append(live, pod)
			}
		}
		if int32(len(live)) != expectedReplicas {
			return false, nil
		}

		observedMilliCPU = 0
		oldestSample = time.Time{}
		for _, pod := range live {
			podMetrics, err := metrics.MetricsV1beta1().PodMetricses(namespace).Get(ctx, pod.Name, metav1.GetOptions{})
			if apierrors.IsNotFound(err) || apierrors.IsServiceUnavailable(err) {
				return false, nil
			}
			if err != nil {
				return false, err
			}
			sampledAt := podMetrics.Timestamp.Time
			if sampledAt.IsZero() || sampledAt.Before(notBefore) || sampledAt.Before(pod.CreationTimestamp.Time) {
				return false, nil
			}
			if oldestSample.IsZero() || sampledAt.Before(oldestSample) {
				oldestSample = sampledAt
			}
			for _, container := range podMetrics.Containers {
				observedMilliCPU += container.Usage.Cpu().MilliValue()
			}
		}
		return observedMilliCPU >= minimumMilliCPU &&
			(maximumMilliCPU < 0 || observedMilliCPU <= maximumMilliCPU), nil
	})
	if err != nil {
		return observedMilliCPU, oldestSample, fmt.Errorf(
			"wait for aggregate TiDB CPU in group %s/%s with %d replica(s) in [%dm,%dm] after %s (last=%dm at %s, Pods=%v): %w",
			namespace, group, expectedReplicas, minimumMilliCPU, maximumMilliCPU,
			notBefore.Format(time.RFC3339), observedMilliCPU, oldestSample.Format(time.RFC3339), lastPods.Items, err,
		)
	}
	return observedMilliCPU, oldestSample, nil
}

// TiDBGroupUnchanged watches the target capacity for the requested duration.
func TiDBGroupUnchanged(
	ctx context.Context,
	c client.Client,
	target *tioperatorv1alpha1.TiDBGroup,
	expectedCPU string,
	expectedReplicas int32,
	expectedTierLabel string,
	duration time.Duration,
) error {
	wantCPU := resource.MustParse(expectedCPU)
	deadline := time.NewTimer(duration)
	defer deadline.Stop()
	ticker := time.NewTicker(PollInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return fmt.Errorf("verify TiDBGroup capacity unchanged: %w", ctx.Err())
		case <-deadline.C:
			return nil
		case <-ticker.C:
			current := &tioperatorv1alpha1.TiDBGroup{}
			if err := c.Get(ctx, client.ObjectKeyFromObject(target), current); err != nil {
				return fmt.Errorf("get Autoscaler target: %w", err)
			}
			cpu := current.Spec.Template.Spec.Resources.CPU
			actualReplicas := ptr.Deref(current.Spec.Replicas, 0)
			actualTierLabel := current.Labels[data.TierLabelKey]
			if actualReplicas != expectedReplicas ||
				actualTierLabel != expectedTierLabel ||
				cpu == nil || cpu.Cmp(wantCPU) != 0 {
				actualCPU := "<nil>"
				if cpu != nil {
					actualCPU = cpu.String()
				}
				return fmt.Errorf(
					"Autoscaler changed target capacity: CPU=%s replicas=%d tierLabel=%q, want CPU=%s replicas=%d tierLabel=%q",
					actualCPU, actualReplicas, actualTierLabel,
					expectedCPU, expectedReplicas, expectedTierLabel,
				)
			}
		}
	}
}

func logf(format string, args ...any) {
	fmt.Fprintf(ginkgo.GinkgoWriter, format+"\n", args...)
}