/*
Copyright 2025 The Volcano Authors.
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.
*/

package sharding

import (
	"context"
	"fmt"
	"sort"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	kubeinformer "k8s.io/client-go/informers"
	kubefake "k8s.io/client-go/kubernetes/fake"
	"k8s.io/klog/v2"
	"sigs.k8s.io/yaml"

	vcfake "volcano.sh/apis/pkg/client/clientset/versioned/fake"
	vcinformer "volcano.sh/apis/pkg/client/informers/externalversions"

	"volcano.sh/volcano/pkg/controllers/framework"
)

// TestShardingController is a test wrapper for ShardingController
type TestShardingController struct {
	Controller *ShardingController
	StopCh     chan struct{}
}

// TestControllerOption contains options for test controller
type TestControllerOption struct {
	InitialObjects   []runtime.Object
	SchedulerConfigs []SchedulerConfigSpec
	ShardSyncPeriod  time.Duration
	StopCh           chan struct{}
	// SkipConfigMap, when true, omits the sharding ConfigMap from the fake
	// client so the controller exercises the flag-based fallback path.
	SkipConfigMap bool
}

// buildShardingConfigMap creates a ConfigMap containing the given scheduler
// configs.
func buildShardingConfigMap(specs []SchedulerConfigSpec) *corev1.ConfigMap {
	cfg := ShardingConfig{SchedulerConfigs: specs}
	data, err := yaml.Marshal(cfg)
	if err != nil {
		panic(fmt.Sprintf("failed to marshal sharding config for test: %v", err))
	}
	return &corev1.ConfigMap{
		ObjectMeta: metav1.ObjectMeta{
			Name:      DefaultConfigMapName,
			Namespace: DefaultConfigMapNamespace,
		},
		Data: map[string]string{
			ConfigMapDataKey: string(data),
		},
	}
}

// newTestShardingController creates a new test controller.
// It mirrors the production flow: scheduler configs are loaded from a
// ConfigMap (created in the fake client) during Run().
func newTestShardingController(t *testing.T, opt *TestControllerOption) *TestShardingController {
	// Create controller
	controller := &ShardingController{
		shardSyncPeriod: defaultShardSyncPeriod,
	}

	if opt == nil {
		opt = &TestControllerOption{}
	}

	// Set default scheduler configs if none provided
	if len(opt.SchedulerConfigs) == 0 {
		klog.Infof("Using default scheduler configs for test")
		opt.SchedulerConfigs = NewShardingControllerOptions().SchedulerConfigs
		klog.Infof("Default scheduler configs: %+v", opt.SchedulerConfigs)
	}

	if opt.StopCh == nil {
		opt.StopCh = make(chan struct{})
	}

	// Build a sharding ConfigMap from the requested scheduler configs and
	// include it in the fake client's initial objects (unless skipped).
	initialObjects := opt.InitialObjects
	if !opt.SkipConfigMap {
		cm := buildShardingConfigMap(opt.SchedulerConfigs)
		initialObjects = append(initialObjects, cm)
	}

	// Create fake clients
	kubeClient := kubefake.NewSimpleClientset(initialObjects...)
	vcClient := vcfake.NewSimpleClientset()

	// Create informer factories
	kubeInformerFactory := kubeinformer.NewSharedInformerFactory(kubeClient, time.Millisecond*500)
	vcInformerFactory := vcinformer.NewSharedInformerFactory(vcClient, time.Millisecond*500)

	// Create controller options
	controllerOpt := &framework.ControllerOption{
		KubeClient:              kubeClient,
		VolcanoClient:           vcClient,
		SharedInformerFactory:   kubeInformerFactory,
		VCSharedInformerFactory: vcInformerFactory,
	}

	// Initialize controller -- configs will be loaded from ConfigMap in Run().
	err := controller.Initialize(controllerOpt)
	assert.NoError(t, err, "should initialize controller successfully")

	// Start controller in separate goroutine
	go controller.Run(opt.StopCh)

	// Wait for controller to start
	time.Sleep(300 * time.Millisecond)

	return &TestShardingController{
		Controller: controller,
		StopCh:     opt.StopCh,
	}
}

func closeTestShardingController(controller *TestShardingController) {
	close(controller.StopCh)
}

// WaitForNodeMetricsUpdate waits for node metrics to be updated
func WaitForNodeMetricsUpdate(t *testing.T, controller *ShardingController, nodeName string, timeout time.Duration) {
	start := time.Now()
	for {
		controller.metricsMutex.RLock()
		metrics := controller.GetNodeMetrics(nodeName)
		controller.metricsMutex.RUnlock()
		if metrics != nil && !metrics.LastUpdated.IsZero() {
			return
		}

		if time.Since(start) > timeout {
			t.Fatalf("Timeout waiting for node metrics to be updated for %s", nodeName)
		}

		time.Sleep(50 * time.Millisecond)
	}
}

// WaitForQueueProcessing waits for queue to be processed
func WaitForQueueProcessing(controller *ShardingController, timeout time.Duration) error {
	start := time.Now()
	for {
		if controller.nodeShardQueue.Len() == 0 && controller.nodeEventQueue.Len() == 0 {
			return nil
		}

		if time.Since(start) > timeout {
			return fmt.Errorf("timeout waiting for queue to empty, main queue: %d, node event queue: %d",
				controller.nodeShardQueue.Len(), controller.nodeEventQueue.Len())
		}

		time.Sleep(10 * time.Millisecond)
	}
}

// ForceSyncShards forces a shard synchronization and waits for completion
func ForceSyncShards(t *testing.T, controller *ShardingController, timeout time.Duration) {
	controller.syncShards()
	assert.NoError(t, WaitForQueueProcessing(controller, timeout), "queue should be processed within timeout")
	// Clean up queue
	for controller.nodeShardQueue.Len() > 0 {
		controller.processNextItem()
	}
}

func fakeShardInUse(t *testing.T, controller *ShardingController, schedulerName string) {
	shard, err := controller.vcClient.ShardV1alpha1().NodeShards().Get(context.TODO(), schedulerName, metav1.GetOptions{})
	assert.NoError(t, err, "should get shard for %s", schedulerName)
	shard.Status.NodesInUse = shard.Spec.NodesDesired
	controller.vcClient.ShardV1alpha1().NodeShards().Update(context.TODO(), shard, metav1.UpdateOptions{})
}

// VerifyAssignment verifies assignment for a scheduler
func VerifyAssignment(t *testing.T, controller *ShardingController, schedulerName string, expectedNodes []string) {
	shard, err := controller.vcClient.ShardV1alpha1().NodeShards().Get(context.TODO(), schedulerName, metav1.GetOptions{})
	assert.NoError(t, err, "should get shard for %s", schedulerName)

	// Sort nodes for consistent comparison
	actualNodes := shard.Spec.NodesDesired
	sort.Strings(actualNodes)
	sort.Strings(expectedNodes)

	assert.ElementsMatch(t, expectedNodes, actualNodes,
		"nodes assigned to %s should match expected", schedulerName)
}

func VerifyAssignmentUpdate(t *testing.T, controller *ShardingController, schedulerName string, expectedNodesToAdd []string, expectedNodesToRemove []string) {
	shard, err := controller.vcClient.ShardV1alpha1().NodeShards().Get(context.TODO(), schedulerName, metav1.GetOptions{})
	assert.NoError(t, err, "should get shard for %s", schedulerName)

	// Sort nodes for consistent comparison
	assert.ElementsMatch(t, shard.Status.NodesToAdd, expectedNodesToAdd, "nodesToAdd assigned to %s should match expected", schedulerName)
	assert.ElementsMatch(t, shard.Status.NodesToRemove, expectedNodesToRemove, "nodesToRemove assigned to %s should match expected", schedulerName)
}

// SetupPodsOnNode creates multiple pods on a specific node
func SetupPodsOnNode(t *testing.T, controller *ShardingController, nodeName string, pods []*corev1.Pod) {
	for _, pod := range pods {
		pod.Spec.NodeName = nodeName
		_, err := controller.kubeClient.CoreV1().Pods(pod.Namespace).Create(context.TODO(), pod, metav1.CreateOptions{})
		assert.NoError(t, err, "should create pod %s", pod.Name)
	}
}

// CreateTestNode creates a test node with specified configuration
func CreateTestNode(name string, cpu string, isWarmup bool, labels map[string]string) *corev1.Node {
	if labels == nil {
		labels = make(map[string]string)
	}

	if isWarmup {
		labels["node.volcano.sh/warmup"] = "true"
	}

	return &corev1.Node{
		ObjectMeta: metav1.ObjectMeta{
			Name:   name,
			Labels: labels,
		},
		Status: corev1.NodeStatus{
			Capacity: corev1.ResourceList{
				corev1.ResourceCPU:    resource.MustParse(cpu),
				corev1.ResourceMemory: resource.MustParse("32Gi"),
				corev1.ResourcePods:   *resource.NewQuantity(110, resource.DecimalExponent),
			},
			Allocatable: corev1.ResourceList{
				corev1.ResourceCPU:    resource.MustParse(cpu),
				corev1.ResourceMemory: resource.MustParse("32Gi"),
				corev1.ResourcePods:   *resource.NewQuantity(110, resource.DecimalExponent),
			},
		},
	}
}

// CreateTestPod creates a test pod with specified resource requests
func CreateTestPod(namespace, name, nodeName, cpu, memory string) *corev1.Pod {
	return &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: namespace,
		},
		Spec: corev1.PodSpec{
			NodeName: nodeName,
			Containers: []corev1.Container{
				{
					Name: "test-container",
					Resources: corev1.ResourceRequirements{
						Requests: corev1.ResourceList{
							corev1.ResourceCPU:    resource.MustParse(cpu),
							corev1.ResourceMemory: resource.MustParse(memory),
						},
					},
				},
			},
		},
		Status: corev1.PodStatus{
			Phase: corev1.PodRunning,
		},
	}
}

// CreateTestNodeWithPods creates a test node with multiple pods
func CreateTestNodeWithPods(t *testing.T, controller *ShardingController, nodeName string, cpu string, isWarmup bool, pods []*corev1.Pod) {
	// Create node
	node := CreateTestNode(nodeName, cpu, isWarmup, nil)
	_, err := controller.kubeClient.CoreV1().Nodes().Create(context.TODO(), node, metav1.CreateOptions{})
	assert.NoError(t, err, "should create node %s", nodeName)

	// Create pods
	for _, pod := range pods {
		pod.Spec.NodeName = nodeName
		_, err := controller.kubeClient.CoreV1().Pods(pod.Namespace).Create(context.TODO(), pod, metav1.CreateOptions{})
		assert.NoError(t, err, "should create pod %s", pod.Name)
	}

	// Wait for cache sync
	time.Sleep(300 * time.Millisecond)
}

// CleanupController cleans up controller resources
func CleanupController(testCtrl *TestShardingController) {
	if testCtrl == nil {
		return
	}

	// First, process any remaining items in queues
	if testCtrl.Controller != nil {
		// Process main queue
		for testCtrl.Controller.nodeShardQueue.Len() > 0 {
			testCtrl.Controller.processNextItem()
		}
		klog.Infof("Processed remaining items in main queue")

		// Process node event queue
		for testCtrl.Controller.nodeEventQueue.Len() > 0 {
			key, quit := testCtrl.Controller.nodeEventQueue.Get()
			if quit {
				break
			}
			testCtrl.Controller.nodeEventQueue.Done(key)
		}
		klog.Infof("Processed remaining items in queues")
	}

	// Close stop channel if exists
	if testCtrl.StopCh != nil {
		select {
		case <-testCtrl.StopCh:
			// Already closed
		default:
			close(testCtrl.StopCh)
		}
		testCtrl.StopCh = nil
	}
	klog.Infof("Closed stop channel")

	// Clear controller reference
	if testCtrl.Controller != nil {
		testCtrl.Controller = nil
	}
}

// getDefaultTestSchedulerConfigs returns default scheduler configs for testing
func getDefaultTestSchedulerConfigs() []SchedulerConfigSpec {
	return []SchedulerConfigSpec{
		{
			Name:              "volcano-scheduler",
			Type:              "volcano",
			CPUUtilizationMin: 0.0,
			CPUUtilizationMax: 0.6,
			PreferWarmupNodes: false,
			MinNodes:          1,
			MaxNodes:          10,
		},
		{
			Name:              "agent-scheduler",
			Type:              "agent",
			CPUUtilizationMin: 0.7,
			CPUUtilizationMax: 1.0,
			PreferWarmupNodes: true,
			MinNodes:          1,
			MaxNodes:          10,
		},
	}
}
