/*
Copyright 2018 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 mutate

import (
	"encoding/json"
	"fmt"
	"strconv"

	admissionv1 "k8s.io/api/admission/v1"
	whv1 "k8s.io/api/admissionregistration/v1"
	v1 "k8s.io/api/core/v1"
	"k8s.io/klog/v2"

	"volcano.sh/apis/pkg/apis/batch/v1alpha1"
	"volcano.sh/volcano/pkg/controllers/job/plugins/distributed-framework/mpi"
	"volcano.sh/volcano/pkg/controllers/job/plugins/distributed-framework/pytorch"
	"volcano.sh/volcano/pkg/controllers/job/plugins/distributed-framework/ray"
	"volcano.sh/volcano/pkg/controllers/job/plugins/distributed-framework/tensorflow"
	commonutil "volcano.sh/volcano/pkg/util"
	"volcano.sh/volcano/pkg/webhooks/router"
	"volcano.sh/volcano/pkg/webhooks/schema"
	"volcano.sh/volcano/pkg/webhooks/util"
)

const (
	// DefaultQueue constant stores the name of the queue as "default"
	DefaultQueue = "default"
	// DefaultMaxRetry is the default number of retries.
	DefaultMaxRetry = 3

	defaultMaxRetry int32 = 3
)

func init() {
	router.RegisterAdmission(service)
}

var service = &router.AdmissionService{
	Path: "/jobs/mutate",
	Func: Jobs,

	Config: config,

	MutatingConfig: &whv1.MutatingWebhookConfiguration{
		Webhooks: []whv1.MutatingWebhook{{
			Name: "mutatejob.volcano.sh",
			Rules: []whv1.RuleWithOperations{
				{
					Operations: []whv1.OperationType{whv1.Create},
					Rule: whv1.Rule{
						APIGroups:   []string{"batch.volcano.sh"},
						APIVersions: []string{"v1alpha1"},
						Resources:   []string{"jobs"},
					},
				},
			},
		}},
	},
}

var config = &router.AdmissionServiceConfig{}

type patchOperation struct {
	Op    string      `json:"op"`
	Path  string      `json:"path"`
	Value interface{} `json:"value,omitempty"`
}

// Jobs mutate jobs.
func Jobs(ar admissionv1.AdmissionReview) *admissionv1.AdmissionResponse {
	klog.V(3).Infof("mutating jobs")

	job, err := schema.DecodeJob(ar.Request.Object, ar.Request.Resource)
	if err != nil {
		return util.ToAdmissionResponse(err)
	}

	var patchBytes []byte
	switch ar.Request.Operation {
	case admissionv1.Create:
		patchBytes, _ = createPatch(job)
	default:
		err = fmt.Errorf("expect operation to be 'CREATE' ")
		return util.ToAdmissionResponse(err)
	}

	klog.V(3).Infof("AdmissionResponse: patch=%v", string(patchBytes))
	reviewResponse := admissionv1.AdmissionResponse{
		Allowed: true,
		Patch:   patchBytes,
	}
	if len(patchBytes) > 0 {
		pt := admissionv1.PatchTypeJSONPatch
		reviewResponse.PatchType = &pt
	}

	return &reviewResponse
}

func createPatch(job *v1alpha1.Job) ([]byte, error) {
	var patch []patchOperation
	pathQueue := patchDefaultQueue(job)
	if pathQueue != nil {
		patch = append(patch, *pathQueue)
	}
	pathScheduler := patchDefaultScheduler(job)
	if pathScheduler != nil {
		patch = append(patch, *pathScheduler)
	}
	pathMaxRetry := patchDefaultMaxRetry(job)
	if pathMaxRetry != nil {
		patch = append(patch, *pathMaxRetry)
	}
	pathSpec := mutateSpec(job.Spec.Tasks, "/spec/tasks", job)
	if pathSpec != nil {
		patch = append(patch, *pathSpec)
	}
	pathMinAvailable := patchDefaultMinAvailable(job)
	if pathMinAvailable != nil {
		patch = append(patch, *pathMinAvailable)
	}
	// Add default plugins for some distributed-framework plugin cases
	patchPlugins := patchDefaultPlugins(job)
	if patchPlugins != nil {
		patch = append(patch, *patchPlugins)
	}
	return json.Marshal(patch)
}

func patchDefaultQueue(job *v1alpha1.Job) *patchOperation {
	//Add default queue if not specified.
	if job.Spec.Queue == "" {
		return &patchOperation{Op: "add", Path: "/spec/queue", Value: DefaultQueue}
	}
	return nil
}

func patchDefaultScheduler(job *v1alpha1.Job) *patchOperation {
	// Add default scheduler name if not specified.
	if job.Spec.SchedulerName == "" {
		return &patchOperation{Op: "add", Path: "/spec/schedulerName", Value: commonutil.GenerateSchedulerName(config.SchedulerNames)}
	}
	return nil
}

func patchDefaultMaxRetry(job *v1alpha1.Job) *patchOperation {
	// Add default maxRetry if maxRetry is zero.
	if job.Spec.MaxRetry == 0 {
		return &patchOperation{Op: "add", Path: "/spec/maxRetry", Value: DefaultMaxRetry}
	}
	return nil
}

func patchDefaultMinAvailable(job *v1alpha1.Job) *patchOperation {
	// Add default minAvailable if minAvailable is zero.
	if job.Spec.MinAvailable == 0 {
		var jobMinAvailable int32
		for _, task := range job.Spec.Tasks {
			if task.MinAvailable != nil {
				jobMinAvailable += *task.MinAvailable
			} else {
				jobMinAvailable += task.Replicas
			}
		}

		return &patchOperation{Op: "add", Path: "/spec/minAvailable", Value: jobMinAvailable}
	}
	return nil
}

func mutateSpec(tasks []v1alpha1.TaskSpec, basePath string, job *v1alpha1.Job) *patchOperation {
	// TODO: Enable this configuration when dependOn supports coexistence with the gang plugin
	// if _, ok := job.Spec.Plugins[mpi.MpiPluginName]; ok {
	// 	mpi.AddDependsOn(job)
	// }
	patched := false
	for index := range tasks {
		// add default task name
		taskName := tasks[index].Name
		if len(taskName) == 0 {
			patched = true
			tasks[index].Name = v1alpha1.DefaultTaskSpec + strconv.Itoa(index)
		}

		if tasks[index].Template.Spec.HostNetwork && tasks[index].Template.Spec.DNSPolicy == "" {
			patched = true
			tasks[index].Template.Spec.DNSPolicy = v1.DNSClusterFirstWithHostNet
		}

		if tasks[index].MinAvailable == nil {
			patched = true
			// Only derive minAvailable from the partition policy when minPartitions is
			// explicitly set (> 0). Otherwise minPartitions defaults to 0, which would
			// make minAvailable 0 and silently drop the task's gang guarantee, so fall
			// back to replicas just like a task without a partition policy.
			if tasks[index].PartitionPolicy != nil && tasks[index].PartitionPolicy.MinPartitions > 0 {
				minAvailable := tasks[index].PartitionPolicy.MinPartitions * tasks[index].PartitionPolicy.PartitionSize
				tasks[index].MinAvailable = &minAvailable
			} else {
				minAvailable := tasks[index].Replicas
				tasks[index].MinAvailable = &minAvailable
			}
		}

		if tasks[index].MaxRetry == 0 {
			patched = true
			tasks[index].MaxRetry = defaultMaxRetry
		}
	}
	if !patched {
		return nil
	}
	return &patchOperation{
		Op:    "replace",
		Path:  basePath,
		Value: tasks,
	}
}

func patchDefaultPlugins(job *v1alpha1.Job) *patchOperation {
	if job.Spec.Plugins == nil {
		return nil
	}
	plugins := map[string][]string{}
	for k, v := range job.Spec.Plugins {
		plugins[k] = v
	}

	// Because the tensorflow-plugin, mpi-plugin and pytorch-plugin depend on svc-plugin.
	// If the svc-plugin is not defined, we should add it.
	_, hasTf := job.Spec.Plugins[tensorflow.TFPluginName]
	_, hasMPI := job.Spec.Plugins[mpi.MPIPluginName]
	_, hasPytorch := job.Spec.Plugins[pytorch.PytorchPluginName]
	_, hasRay := job.Spec.Plugins[ray.RayPluginName]
	if hasTf || hasMPI || hasPytorch || hasRay {
		if _, ok := plugins["svc"]; !ok {
			plugins["svc"] = []string{}
		}
	}

	if _, ok := job.Spec.Plugins["mpi"]; ok {
		if _, ok := plugins["ssh"]; !ok {
			plugins["ssh"] = []string{}
		}
	}

	return &patchOperation{
		Op:    "replace",
		Path:  "/spec/plugins",
		Value: plugins,
	}
}
