package k8s

import (
	"context"
	_ "embed"
	"fmt"

	"github.com/imdario/mergo"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/runtime/schema"

	"github.com/pingcap/test-infra/sdk/resource"
)

var (
	//go:embed assets/schedule.yaml
	scheduleSpec        []byte
	defaultScheduleSpec map[string]interface{}
	schedulesGVR        = schema.GroupVersionResource{Group: "chaos-mesh.org", Version: "v1alpha1", Resource: "schedules"}
)

func init() {
	mustUnmarshalYaml(scheduleSpec, &defaultScheduleSpec)
}

type ChaosClientV2 struct {
	KubeClientSet
}

func (c *ChaosClientV2) Create(opts resource.InjectChaosOptions) (resource.ChaosIdentifier, error) {
	switch opts.Kind() {
	case resource.ExperimentNetworkDelayKind, resource.ExperimentNetworkPartitionKind, resource.ExperimentNetworkLossKind:
		if opts.Target == nil || len(*opts.Target) == 0 {
			sugar.Warnf("the target field needs to be set because the tidbcloud cluster uses the host network, otherwise it will affect the network whole scope of the host")
		}
	}
	switch opts.Kind() {
	case resource.ExperimentFailureKind:
		return c.createPodFailureObject(opts, opts.ExperimentSpec.(*resource.FailureExperimentSpec))
	case resource.ExperimentNetworkDelayKind:
		return c.createNetworkDelayObject(opts, opts.ExperimentSpec.(*resource.NetworkLatencySpec))
	case resource.ExperimentNetworkPartitionKind:
		return c.createNetworkPartitionObject(opts, opts.ExperimentSpec.(*resource.NetworkPartitionSpec))
	case resource.ExperimentNetworkLossKind:
		return c.createNetworkLossObject(opts, opts.ExperimentSpec.(*resource.NetworkLossSpec))
	case resource.ExperimentIODelayKind:
		return c.createIODelayObject(opts, opts.ExperimentSpec.(*resource.IODelaySpec))
	case resource.ExperimentTimeChaosKind:
		return c.createTimeChaosObject(opts, opts.ExperimentSpec.(*resource.TimeChaosSpec))
	case resource.ExperimentStressChaosKind:
		return c.createStressChaosObject(opts, opts.ExperimentSpec.(*resource.StressChaosSpec))
	}
	return nil, fmt.Errorf("unknown chaos kind: %s", opts.Kind())
}

func (c *ChaosClientV2) Delete(id resource.ChaosIdentifier) error {
	chaosId := id.(*ChaosIdentifier)
	switch chaosId.Spec.Kind() {
	case resource.ExperimentFailureKind:
		return c.DynamicClient.Resource(podChaosGVR).Namespace(c.Namespace).Delete(context.TODO(), chaosId.Name, metav1.DeleteOptions{})
	case resource.ExperimentNetworkDelayKind:
		return c.DynamicClient.Resource(networkChaosGVR).Namespace(c.Namespace).Delete(context.TODO(), chaosId.Name, metav1.DeleteOptions{})
	case resource.ExperimentNetworkPartitionKind:
		return c.DynamicClient.Resource(networkChaosGVR).Namespace(c.Namespace).Delete(context.TODO(), chaosId.Name, metav1.DeleteOptions{})
	case resource.ExperimentNetworkLossKind:
		return c.DynamicClient.Resource(networkChaosGVR).Namespace(c.Namespace).Delete(context.TODO(), chaosId.Name, metav1.DeleteOptions{})
	case resource.ExperimentIODelayKind:
		return c.DynamicClient.Resource(ioChaosGVR).Namespace(c.Namespace).Delete(context.TODO(), chaosId.Name, metav1.DeleteOptions{})
	case resource.ExperimentTimeChaosKind:
		return c.DynamicClient.Resource(timeChaosGVR).Namespace(c.Namespace).Delete(context.TODO(), chaosId.Name, metav1.DeleteOptions{})
	case resource.ExperimentStressChaosKind:
		return c.DynamicClient.Resource(stressChaosGVR).Namespace(c.Namespace).Delete(context.TODO(), chaosId.Name, metav1.DeleteOptions{})
	case resource.ExperimentSchedulesKind:
		return c.DynamicClient.Resource(schedulesGVR).Namespace(c.Namespace).Delete(context.TODO(), chaosId.Name, metav1.DeleteOptions{})
	}
	return fmt.Errorf("unknown chaos identifier kind: %s", chaosId.Spec.Kind())
}

func (c *ChaosClientV2) createPodFailureObject(opts resource.InjectChaosOptions, spec *resource.FailureExperimentSpec) (*ChaosIdentifier, error) {
	var cis resource.ExperimentSpec = spec
	gvr := podChaosGVR
	if len(opts.Selector) == 0 {
		return nil, fmt.Errorf("expect a non-empty pod selector")
	}
	if opts.Target != nil {
		return nil, fmt.Errorf("the opts.Target should be nil when create a failure injection")
	}

	failureObj := new(unstructured.Unstructured)
	if err := mergo.Merge(&failureObj.Object, runtime.DeepCopyJSON(defaultPodFailureSpec)); err != nil {
		return nil, err
	}
	failureObj.SetName(resource.GenResourceName("", "pod-failure-"))

	podSelector := c.buildPodSelector(opts.Selector)
	_ = unstructured.SetNestedField(failureObj.Object, podSelector.Pods, "spec", "selector", "pods")
	if spec.Duration != "" {
		_ = unstructured.SetNestedField(failureObj.Object, spec.Duration, "spec", "duration")
	}
	if spec.Scheduler != nil {
		var err error
		failureObj, err = c.wrapScheduleObject(failureObj, spec.Scheduler)
		if err != nil {
			return nil, err
		}
		gvr = schedulesGVR
		cis = &resource.SchedulerExperimentSpecWrapper{ExperimentSpec: spec}
	}
	obj, err := c.DynamicClient.Resource(gvr).Namespace(c.Namespace).Create(context.TODO(), failureObj, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	return &ChaosIdentifier{Namespace: c.Namespace, Name: obj.GetName(), Spec: cis}, nil
}

func (c *ChaosClientV2) createNetworkPartitionObject(opts resource.InjectChaosOptions, spec *resource.NetworkPartitionSpec) (resource.ChaosIdentifier, error) {
	var cis resource.ExperimentSpec = spec
	gvr := networkChaosGVR
	if len(opts.Selector) == 0 {
		return nil, fmt.Errorf("expect a non-empty pod selector")
	}
	if opts.Target == nil || len(*opts.Target) == 0 {
		return nil, fmt.Errorf("expect a non-empty target pod selector")
	}

	partitionObj := new(unstructured.Unstructured)
	if err := mergo.Merge(&partitionObj.Object, runtime.DeepCopyJSON(defaultNetworkPartitionSpec)); err != nil {
		return nil, err
	}
	partitionObj.SetName(resource.GenResourceName("", "network-partition-"))

	podSelector := c.buildPodSelector(opts.Selector)
	targetPodSelector := c.buildPodSelector(*opts.Target)
	_ = unstructured.SetNestedField(partitionObj.Object, podSelector.Pods, "spec", "selector", "pods")
	_ = unstructured.SetNestedField(partitionObj.Object, targetPodSelector.Pods, "spec", "target", "selector", "pods")
	if spec.Duration != "" {
		_ = unstructured.SetNestedField(partitionObj.Object, spec.Duration, "spec", "duration")
	}
	if spec.Direction != "" {
		_ = unstructured.SetNestedField(partitionObj.Object, spec.Direction, "spec", "direction")
	}
	if spec.Scheduler != nil {
		var err error
		partitionObj, err = c.wrapScheduleObject(partitionObj, spec.Scheduler)
		if err != nil {
			return nil, err
		}
		gvr = schedulesGVR
		cis = &resource.SchedulerExperimentSpecWrapper{ExperimentSpec: spec}
	}

	obj, err := c.DynamicClient.Resource(gvr).Namespace(c.Namespace).Create(context.TODO(), partitionObj, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	return &ChaosIdentifier{Namespace: c.Namespace, Name: obj.GetName(), Spec: cis}, nil
}

func (c *ChaosClientV2) createNetworkLossObject(opts resource.InjectChaosOptions, spec *resource.NetworkLossSpec) (resource.ChaosIdentifier, error) {
	var cis resource.ExperimentSpec = spec
	gvr := networkChaosGVR
	if len(opts.Selector) == 0 {
		return nil, fmt.Errorf("expect a non-empty pod selector")
	}
	lossObj := new(unstructured.Unstructured)
	if err := mergo.Merge(&lossObj.Object, runtime.DeepCopyJSON(defaultNetworkLossSpec)); err != nil {
		return nil, err
	}
	lossObj.SetName(resource.GenResourceName("", "network-loss-"))

	podSelector := c.buildPodSelector(opts.Selector)
	_ = unstructured.SetNestedField(lossObj.Object, podSelector.Pods, "spec", "selector", "pods")
	if opts.Target != nil && len(*opts.Target) > 0 {
		targetPodSelector := c.buildPodSelector(*opts.Target)
		_ = unstructured.SetNestedField(lossObj.Object, "to", "spec", "direction")
		_ = unstructured.SetNestedField(lossObj.Object, targetPodSelector.Pods, "spec", "target", "selector", "pods")
		_ = unstructured.SetNestedField(lossObj.Object, "all", "spec", "target", "mode")
	}
	if spec.Duration != "" {
		_ = unstructured.SetNestedField(lossObj.Object, spec.Duration, "spec", "duration")
	}
	_ = unstructured.SetNestedField(lossObj.Object, spec.Loss, "spec", "loss", "loss")
	_ = unstructured.SetNestedField(lossObj.Object, spec.Correlation, "spec", "loss", "correlation")

	if spec.Scheduler != nil {
		var err error
		lossObj, err = c.wrapScheduleObject(lossObj, spec.Scheduler)
		if err != nil {
			return nil, err
		}
		gvr = schedulesGVR
		cis = &resource.SchedulerExperimentSpecWrapper{ExperimentSpec: spec}
	}
	obj, err := c.DynamicClient.Resource(gvr).Namespace(c.Namespace).Create(context.TODO(), lossObj, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}

	return &ChaosIdentifier{Namespace: c.Namespace, Name: obj.GetName(), Spec: cis}, nil
}

func (c *ChaosClientV2) createNetworkDelayObject(opts resource.InjectChaosOptions, spec *resource.NetworkLatencySpec) (resource.ChaosIdentifier, error) {
	var cis resource.ExperimentSpec = spec
	gvr := networkChaosGVR
	if len(opts.Selector) == 0 {
		return nil, fmt.Errorf("expect a non-empty pod selector")
	}
	delayObj := new(unstructured.Unstructured)
	if err := mergo.Merge(&delayObj.Object, runtime.DeepCopyJSON(defaultNetworkDelaySpec)); err != nil {
		return nil, err
	}
	delayObj.SetName(resource.GenResourceName("", "network-delay-"))

	podSelector := c.buildPodSelector(opts.Selector)
	_ = unstructured.SetNestedField(delayObj.Object, podSelector.Pods, "spec", "selector", "pods")
	if opts.Target != nil && len(*opts.Target) > 0 {
		targetPodSelector := c.buildPodSelector(*opts.Target)
		_ = unstructured.SetNestedField(delayObj.Object, "to", "spec", "direction")
		_ = unstructured.SetNestedField(delayObj.Object, targetPodSelector.Pods, "spec", "target", "selector", "pods")
		_ = unstructured.SetNestedField(delayObj.Object, "all", "spec", "target", "mode")
	}
	if spec.Duration != "" {
		_ = unstructured.SetNestedField(delayObj.Object, spec.Duration, "spec", "duration")
	}
	_ = unstructured.SetNestedField(delayObj.Object, spec.Latency, "spec", "delay", "latency")
	_ = unstructured.SetNestedField(delayObj.Object, spec.Jitter, "spec", "delay", "jitter")
	_ = unstructured.SetNestedField(delayObj.Object, spec.Correlation, "spec", "delay", "correlation")
	if spec.Scheduler != nil {
		var err error
		delayObj, err = c.wrapScheduleObject(delayObj, spec.Scheduler)
		if err != nil {
			return nil, err
		}
		gvr = schedulesGVR
		cis = &resource.SchedulerExperimentSpecWrapper{ExperimentSpec: spec}
	}
	obj, err := c.DynamicClient.Resource(gvr).Namespace(c.Namespace).Create(context.TODO(), delayObj, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	return &ChaosIdentifier{Namespace: c.Namespace, Name: obj.GetName(), Spec: cis}, nil
}

func (c *ChaosClientV2) createIODelayObject(opts resource.InjectChaosOptions, spec *resource.IODelaySpec) (resource.ChaosIdentifier, error) {
	var cis resource.ExperimentSpec = spec
	gvr := ioChaosGVR
	if len(opts.Selector) == 0 {
		return nil, fmt.Errorf("expect a non-empty pod selector")
	}
	if opts.Target != nil {
		return nil, fmt.Errorf("the opts.Target should be nil when create a io delay injection")
	}

	delayObj := new(unstructured.Unstructured)
	if err := mergo.Merge(&delayObj.Object, runtime.DeepCopyJSON(defaultIODelaySpec)); err != nil {
		return nil, err
	}
	delayObj.SetName(resource.GenResourceName("", "io-delay-"))

	podSelector := c.buildPodSelector(opts.Selector)
	if err := unstructured.SetNestedField(delayObj.Object, podSelector.Pods, "spec", "selector", "pods"); err != nil {
		return nil, err
	}
	if spec.Duration != "" {
		if err := unstructured.SetNestedField(delayObj.Object, spec.Duration, "spec", "duration"); err != nil {
			return nil, err
		}
	}
	if err := unstructured.SetNestedField(delayObj.Object, spec.Delay, "spec", "delay"); err != nil {
		return nil, err
	}
	if err := unstructured.SetNestedField(delayObj.Object, spec.Percent, "spec", "percent"); err != nil {
		return nil, err
	}
	if err := unstructured.SetNestedField(delayObj.Object, spec.VolumePath, "spec", "volumePath"); err != nil {
		return nil, err
	}
	if err := unstructured.SetNestedField(delayObj.Object, spec.Path, "spec", "path"); err != nil {
		return nil, err
	}

	if spec.Scheduler != nil {
		var err error
		delayObj, err = c.wrapScheduleObject(delayObj, spec.Scheduler)
		if err != nil {
			return nil, err
		}
		gvr = schedulesGVR
		cis = &resource.SchedulerExperimentSpecWrapper{ExperimentSpec: spec}
	}
	obj, err := c.DynamicClient.Resource(gvr).Namespace(c.Namespace).Create(context.TODO(), delayObj, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	return &ChaosIdentifier{Namespace: c.Namespace, Name: obj.GetName(), Spec: cis}, nil
}

func (c *ChaosClientV2) createTimeChaosObject(opts resource.InjectChaosOptions, spec *resource.TimeChaosSpec) (resource.ChaosIdentifier, error) {
	var cis resource.ExperimentSpec = spec
	gvr := timeChaosGVR
	if len(opts.Selector) == 0 {
		return nil, fmt.Errorf("expect a non-empty pod selector")
	}
	if opts.Target != nil {
		return nil, fmt.Errorf("the opts.Target should be nil when create a time chaos injection")
	}

	timeChaosObj := new(unstructured.Unstructured)
	if err := mergo.Merge(&timeChaosObj.Object, runtime.DeepCopyJSON(defaultTimeChaosSpec)); err != nil {
		return nil, err
	}
	timeChaosObj.SetName(resource.GenResourceName("", "time-chaos-"))

	podSelector := c.buildPodSelector(opts.Selector)
	if err := unstructured.SetNestedField(timeChaosObj.Object, podSelector.Pods, "spec", "selector", "pods"); err != nil {
		return nil, err
	}
	if spec.Duration != "" {
		if err := unstructured.SetNestedField(timeChaosObj.Object, spec.Duration, "spec", "duration"); err != nil {
			return nil, err
		}
	}
	if spec.Scheduler != nil {
		var err error
		timeChaosObj, err = c.wrapScheduleObject(timeChaosObj, spec.Scheduler)
		if err != nil {
			return nil, err
		}
		gvr = schedulesGVR
		cis = &resource.SchedulerExperimentSpecWrapper{ExperimentSpec: spec}
	}
	if err := unstructured.SetNestedField(timeChaosObj.Object, spec.TimeOffset, "spec", "timeOffset"); err != nil {
		return nil, err
	}
	obj, err := c.DynamicClient.Resource(gvr).Namespace(c.Namespace).Create(context.TODO(), timeChaosObj, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	return &ChaosIdentifier{Namespace: c.Namespace, Name: obj.GetName(), Spec: cis}, nil
}

func (c *ChaosClientV2) createStressChaosObject(opts resource.InjectChaosOptions, spec *resource.StressChaosSpec) (resource.ChaosIdentifier, error) {
	var cis resource.ExperimentSpec = spec
	gvr := stressChaosGVR
	if len(opts.Selector) == 0 {
		return nil, fmt.Errorf("expect a non-empty pod selector")
	}
	if opts.Target != nil {
		return nil, fmt.Errorf("the opts.Target should be nil when create a stress chaos injection")
	}

	stressObj := new(unstructured.Unstructured)
	if err := mergo.Merge(&stressObj.Object, runtime.DeepCopyJSON(defaultStressChaosSpec)); err != nil {
		return nil, err
	}
	stressObj.SetName(resource.GenResourceName("", "stress-chaos-"))

	podSelector := c.buildPodSelector(opts.Selector)
	if err := unstructured.SetNestedField(stressObj.Object, podSelector.Pods, "spec", "selector", "pods"); err != nil {
		return nil, err
	}
	if spec.Duration != "" {
		if err := unstructured.SetNestedField(stressObj.Object, spec.Duration, "spec", "duration"); err != nil {
			return nil, err
		}
	}
	if spec.Stressors.MemoryStressor != nil {
		if err := unstructured.SetNestedField(stressObj.Object, spec.Stressors.MemoryStressor.Workers, "spec", "stressors", "memory", "workers"); err != nil {
			return nil, err
		}
		if err := unstructured.SetNestedField(stressObj.Object, spec.Stressors.MemoryStressor.Size, "spec", "stressors", "memory", "size"); err != nil {
			return nil, err
		}
	} else {
		if err := unstructured.SetNestedField(stressObj.Object, nil, "spec", "stressors", "memory"); err != nil {
			return nil, err
		}
	}

	if spec.Stressors.CPUStressor != nil {
		if err := unstructured.SetNestedField(stressObj.Object, spec.Stressors.CPUStressor.Workers, "spec", "stressors", "cpu", "workers"); err != nil {
			return nil, err
		}
		if err := unstructured.SetNestedField(stressObj.Object, spec.Stressors.CPUStressor.Load, "spec", "stressors", "cpu", "load"); err != nil {
			return nil, err
		}
	} else {
		if err := unstructured.SetNestedField(stressObj.Object, nil, "spec", "stressors", "cpu"); err != nil {
			return nil, err
		}
	}
	if spec.Scheduler != nil {
		var err error
		stressObj, err = c.wrapScheduleObject(stressObj, spec.Scheduler)
		if err != nil {
			return nil, err
		}
		gvr = schedulesGVR
		cis = &resource.SchedulerExperimentSpecWrapper{ExperimentSpec: spec}
	}
	obj, err := c.DynamicClient.Resource(gvr).Namespace(c.Namespace).Create(context.TODO(), stressObj, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	return &ChaosIdentifier{Namespace: c.Namespace, Name: obj.GetName(), Spec: cis}, nil
}

func (c *ChaosClientV2) buildPodSelector(selector []resource.ChaosSelector) podSelector {
	pods := map[string]interface{}{}
	for _, s := range selector {
		podSelector := s.(*PodSelector)
		list, e := pods[podSelector.Namespace].([]interface{})
		if !e {
			list = []interface{}{}
		}
		list = append(list, podSelector.Name)
		pods[podSelector.Namespace] = list
	}
	return podSelector{pods}
}

func (c *ChaosClientV2) wrapScheduleObject(bodyObject *unstructured.Unstructured, scheduler *resource.ExperimentScheduler) (*unstructured.Unstructured, error) {
	scheduleObj := new(unstructured.Unstructured)
	if err := mergo.Merge(&scheduleObj.Object, runtime.DeepCopyJSON(defaultScheduleSpec)); err != nil {
		return nil, err
	}
	scheduleObj.SetName(bodyObject.GetName())
	_ = unstructured.SetNestedField(scheduleObj.Object, scheduler.Cron, "spec", "schedule")
	_ = unstructured.SetNestedField(scheduleObj.Object, bodyObject.GetKind(), "spec", "type")
	types := map[string]string{
		"PodChaos":     "podChaos",
		"NetworkChaos": "networkChaos",
		"IOChaos":      "ioChaos",
		"TimeChaos":    "timeChaos",
		"StressChaos":  "stressChaos",
	}
	spec, _, _ := unstructured.NestedFieldNoCopy(bodyObject.Object, "spec")
	_ = unstructured.SetNestedField(scheduleObj.Object, spec, "spec", types[bodyObject.GetKind()])
	return scheduleObj, nil
}
