// 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 workload manages SQL load generators used by the real-cluster e2e
// suite.
package workload

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

	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"

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

// SQLConfig describes one continuously running SQL workload.
type SQLConfig struct {
	Name        string
	Namespace   string
	Image       string
	Host        string
	Concurrency int
	Iterations  int
	Period      time.Duration
}

// SQL runs concurrent MySQL clients that execute BENCHMARK queries against
// TiDB. It owns only its workload Pod; the TiDB cluster remains suite-owned.
type SQL struct {
	client client.Client
	pod    *corev1.Pod
}

// NewSQL validates the workload and constructs its Pod without creating it.
func NewSQL(c client.Client, cfg SQLConfig) (*SQL, error) {
	if c == nil {
		return nil, fmt.Errorf("sql workload: Kubernetes client is required")
	}
	if cfg.Name == "" || cfg.Namespace == "" || cfg.Image == "" || cfg.Host == "" {
		return nil, fmt.Errorf("sql workload: name, namespace, image and host are required")
	}
	if cfg.Concurrency < 1 {
		return nil, fmt.Errorf("sql workload: concurrency must be greater than zero")
	}
	if cfg.Iterations < 1 {
		return nil, fmt.Errorf("sql workload: iterations must be greater than zero")
	}
	if cfg.Period < 0 {
		return nil, fmt.Errorf("sql workload: period must not be negative")
	}
	if cfg.Period > 0 && cfg.Period < time.Millisecond {
		return nil, fmt.Errorf("sql workload: period must be at least one millisecond")
	}

	return &SQL{client: c, pod: sqlPod(cfg)}, nil
}

// Start creates the SQL workload and waits until its MySQL readiness query
// succeeds against TiDB.
func (w *SQL) Start(ctx context.Context) error {
	if err := w.client.Create(ctx, w.pod); err != nil {
		return fmt.Errorf("create SQL workload %s/%s: %w", w.pod.Namespace, w.pod.Name, err)
	}

	ready, err := waiter.PodReady(ctx, w.client, w.pod)
	if err != nil {
		return fmt.Errorf("wait for SQL workload %s/%s Ready: %w", w.pod.Namespace, w.pod.Name, err)
	}
	w.pod = ready
	return nil
}

// Stop deletes the workload Pod and waits for it to disappear. It is
// idempotent so suite cleanup can call it after a partially completed test.
func (w *SQL) Stop(ctx context.Context) error {
	err := w.client.Delete(ctx, &corev1.Pod{ObjectMeta: metav1.ObjectMeta{
		Name: w.pod.Name, Namespace: w.pod.Namespace,
	}}, client.GracePeriodSeconds(0))
	if err != nil && !apierrors.IsNotFound(err) {
		return fmt.Errorf("delete SQL workload %s/%s: %w", w.pod.Namespace, w.pod.Name, err)
	}
	if err := waiter.ObjectDeleted(ctx, w.client, w.pod); err != nil {
		return fmt.Errorf("wait for SQL workload %s/%s deletion: %w", w.pod.Namespace, w.pod.Name, err)
	}
	return nil
}

// Pod returns a copy of the workload Pod.
func (w *SQL) Pod() *corev1.Pod {
	return w.pod.DeepCopy()
}

func sqlPod(cfg SQLConfig) *corev1.Pod {
	const mysqlArgs = `--protocol=TCP --ssl-mode=DISABLED --connect-timeout=2 --host="${TIDB_HOST}" --port=4000 --user=root`
	script := `
until mysql ` + mysqlArgs + ` --execute='SELECT 1' >/dev/null 2>&1; do
  sleep 1
done
run_worker() {
  while true; do
    started_at_ms=$(date +%s%3N)
    mysql ` + mysqlArgs + ` --batch --skip-column-names \
      --execute="SELECT BENCHMARK(${ITERATIONS}, SHA2('autoscaling-e2e', 256))" >/dev/null 2>&1 || sleep 1
    if [ "${PERIOD_MILLISECONDS}" -gt "0" ]; then
      finished_at_ms=$(date +%s%3N)
      remaining_ms=$((PERIOD_MILLISECONDS - finished_at_ms + started_at_ms))
      if [ "${remaining_ms}" -gt "0" ]; then
        printf -v remaining_seconds '%d.%03d' "$((remaining_ms / 1000))" "$((remaining_ms % 1000))"
        sleep "${remaining_seconds}"
      fi
    fi
  done
}
worker=0
while [ "${worker}" -lt "${CONCURRENCY}" ]; do
  run_worker &
  worker=$((worker + 1))
done
wait
`
	return &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      cfg.Name,
			Namespace: cfg.Namespace,
			Labels: map[string]string{
				"e2e.autoscaling.pingcap.com/role": "sql-workload",
			},
		},
		Spec: corev1.PodSpec{
			RestartPolicy: corev1.RestartPolicyAlways,
			Containers: []corev1.Container{{
				Name:    "mysql",
				Image:   cfg.Image,
				Command: []string{"/bin/bash", "-ec", script},
				Env: []corev1.EnvVar{
					{Name: "TIDB_HOST", Value: cfg.Host},
					{Name: "CONCURRENCY", Value: strconv.Itoa(cfg.Concurrency)},
					{Name: "ITERATIONS", Value: strconv.Itoa(cfg.Iterations)},
					{Name: "PERIOD_MILLISECONDS", Value: strconv.FormatInt(cfg.Period.Milliseconds(), 10)},
				},
				ReadinessProbe: &corev1.Probe{
					ProbeHandler: corev1.ProbeHandler{Exec: &corev1.ExecAction{Command: []string{
						"/bin/bash", "-ec", "mysql " + mysqlArgs + " --execute='SELECT 1' >/dev/null 2>&1",
					}}},
					InitialDelaySeconds: 1,
					PeriodSeconds:       2,
					TimeoutSeconds:      2,
				},
				Resources: corev1.ResourceRequirements{
					Requests: corev1.ResourceList{
						corev1.ResourceCPU:    resource.MustParse("25m"),
						corev1.ResourceMemory: resource.MustParse("32Mi"),
					},
					Limits: corev1.ResourceList{
						corev1.ResourceCPU:    resource.MustParse("500m"),
						corev1.ResourceMemory: resource.MustParse("256Mi"),
					},
				},
				SecurityContext: &corev1.SecurityContext{
					AllowPrivilegeEscalation: ptr.To(false),
					Capabilities:             &corev1.Capabilities{Drop: []corev1.Capability{"ALL"}},
				},
			}},
		},
	}
}