package k8s

import (
	"context"
	_ "embed"
	"encoding/json"
	"fmt"
	"net/url"
	"os"

	"github.com/creasty/defaults"
	"github.com/hashicorp/go-multierror"
	"github.com/pingcap/errors"
	v1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"
	"sigs.k8s.io/yaml"

	apiresource "github.com/pingcap/test-infra/common/model/resource"
	"github.com/pingcap/test-infra/sdk/core"
	sdkresource "github.com/pingcap/test-infra/sdk/resource"
)

var (
	_ sdkresource.Client = KafkaClusterClient{}

	//go:embed assets/kafka-zk.yaml
	zookeeperSpec []byte
	//go:embed assets/kafka.yaml
	kafkaClusterSpec []byte
)

type KafkaClusterClient struct {
	*KubeClientSet
}

func NewKafkaClusterClient(k *KubeClientSet) KafkaClusterClient {
	return KafkaClusterClient{k}
}

func (t KafkaClusterClient) Create(spec json.RawMessage) (core.Resource, error) {
	sts_zk, sts_kafka, svc_zk_peer, svc_zk, svc_kafka_peer, svc_kafka, apiRes, err := t.buildKafkaClusterObject(spec)
	if err != nil {
		return nil, err
	}
	// if apiRes.Qos != nil {
	// 	if err := t.applyQos(apiRes, sts); err != nil {
	// 		return nil, err
	// 	}
	// }

	svc_zk_peer, err = t.ClientSet.CoreV1().Services(t.Namespace).Create(context.TODO(), svc_zk_peer, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	svc_zk, err = t.ClientSet.CoreV1().Services(t.Namespace).Create(context.TODO(), svc_zk, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	svc_kafka_peer, err = t.ClientSet.CoreV1().Services(t.Namespace).Create(context.TODO(), svc_kafka_peer, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	svc_kafka, err = t.ClientSet.CoreV1().Services(t.Namespace).Create(context.TODO(), svc_kafka, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	sts_zk, err = t.ClientSet.AppsV1().StatefulSets(t.Namespace).Create(context.TODO(), sts_zk, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	sts_kafka, err = t.ClientSet.AppsV1().StatefulSets(t.Namespace).Create(context.TODO(), sts_kafka, metav1.CreateOptions{})
	if err != nil {
		return nil, err
	}
	resource := &KafkaCluster{
		KubeClientSet:      *t.KubeClientSet,
		KafkaCluster:       *apiRes,
		sts_zk:             sts_zk,
		sts_kafka:          sts_kafka,
		service_zk:         svc_zk,
		service_zk_peer:    svc_zk_peer,
		service_kafka:      svc_kafka,
		service_kafka_peer: svc_kafka_peer,
	}
	resource.Provider = t.FQDN()
	if err := resource.Sync(); err != nil {
		return nil, err
	}
	return resource, nil
}

func (t KafkaClusterClient) Get(handle apiresource.ResourceHandle) (core.Resource, error) {
	resource := &KafkaCluster{
		KubeClientSet: *t.KubeClientSet,
		KafkaCluster: apiresource.KafkaCluster{
			ResourceHandle: handle,
		},
	}
	resource.Provider = t.FQDN()
	if err := resource.Sync(); err != nil {
		return nil, err
	}
	return resource, nil
}

func (t KafkaClusterClient) GetEndpoints(_ apiresource.ResourceHandle) (apiresource.ResourceEndpoints, error) {
	return apiresource.ResourceEndpoints{}, nil
}

func (t KafkaClusterClient) Delete(handle apiresource.ResourceHandle) error {
	errs := multierror.Append(nil, t.ClientSet.AppsV1().StatefulSets(t.Namespace).Delete(context.TODO(), handle.Name+"-zk", metav1.DeleteOptions{
		PropagationPolicy: &[]metav1.DeletionPropagation{metav1.DeletePropagationForeground}[0],
	}))
	errs = multierror.Append(errs, t.ClientSet.AppsV1().StatefulSets(t.Namespace).Delete(context.TODO(), handle.Name+"-kafka", metav1.DeleteOptions{
		PropagationPolicy: &[]metav1.DeletionPropagation{metav1.DeletePropagationForeground}[0],
	}))
	errs = multierror.Append(errs, t.ClientSet.CoreV1().Services(t.Namespace).Delete(context.TODO(), handle.Name+"-zk", metav1.DeleteOptions{}))
	errs = multierror.Append(errs, t.ClientSet.CoreV1().Services(t.Namespace).Delete(context.TODO(), handle.Name+"-zk-peer", metav1.DeleteOptions{}))
	errs = multierror.Append(errs, t.ClientSet.CoreV1().Services(t.Namespace).Delete(context.TODO(), handle.Name+"-kafka", metav1.DeleteOptions{}))
	errs = multierror.Append(errs, t.ClientSet.CoreV1().Services(t.Namespace).Delete(context.TODO(), handle.Name+"-kafka-peer", metav1.DeleteOptions{}))
	errs = multierror.Append(errs, t.ClientSet.CoreV1().PersistentVolumeClaims(t.Namespace).DeleteCollection(context.TODO(), metav1.DeleteOptions{}, metav1.ListOptions{
		LabelSelector: labels.Set{TestInfraName: handle.Name}.AsSelector().String(),
	}))
	return errs.ErrorOrNil()
}

func (t KafkaClusterClient) ReserveEnv(handle apiresource.ResourceHandle, action apiresource.GcReserveEnvAction) error {
	return nil
}

func (t KafkaClusterClient) Artifact(handle apiresource.ResourceHandle, action apiresource.GcArtifactAction, archiver sdkresource.ArchiverFunc) (map[string]*url.URL, error) {
	return nil, nil
}

func (t KafkaClusterClient) FQDN() string {
	return sdkresource.FQDN(t)
}

func (t KafkaClusterClient) buildKafkaClusterObject(params json.RawMessage) (*v1.StatefulSet, *v1.StatefulSet, *corev1.Service, *corev1.Service, *corev1.Service, *corev1.Service, *apiresource.KafkaCluster, error) {
	resHandle := &apiresource.KafkaCluster{}
	// TODO: add validation on KafkaCluster
	if err := json.Unmarshal(params, resHandle); err != nil {
		return nil, nil, nil, nil, nil, nil, nil, errors.Trace(err)
	}

	if err := defaults.Set(resHandle); err != nil {
		return nil, nil, nil, nil, nil, nil, nil, errors.Trace(err)
	}

	if resHandle.Spec.KafkaSpec.Env == nil {
		resHandle.Spec.KafkaSpec.Env = make(map[string]string)
	}
	resHandle.Spec.KafkaSpec.Env["ZK_CONNECT"] = resHandle.Name + "-zk." + resHandle.Testbed + ".svc.cluster.local:2181"

	// modify sts_zk
	sts_zk, err := t.buildSTS_ZKObj(resHandle)
	if err != nil {
		return nil, nil, nil, nil, nil, nil, nil, err
	}

	// modify sts_kafka
	sts_kafka, err := t.buildSTS_KafkaObj(resHandle)
	if err != nil {
		return nil, nil, nil, nil, nil, nil, nil, err
	}

	owner, err := GetOwnerReferenceFromTestbed(t.KubeClientSet, resHandle.Testbed)
	if err != nil {
		return nil, nil, nil, nil, nil, nil, nil, err
	}
	svc_zk_peer, svc_zk := t.buildZKService(resHandle.Name, owner)
	svc_kafka_peer, svc_kafka := t.buildKafkaService(resHandle.Name, owner)

	sts_zk.Spec.ServiceName = svc_zk_peer.Name
	sts_kafka.Spec.ServiceName = svc_kafka_peer.Name

	return sts_zk, sts_kafka, svc_zk_peer, svc_zk, svc_kafka_peer, svc_kafka, resHandle, nil
}

func (t KafkaClusterClient) buildSTS_ZKObj(resHandle *apiresource.KafkaCluster) (*v1.StatefulSet, error) {
	sts_zk := &v1.StatefulSet{}
	if err := yaml.Unmarshal(zookeeperSpec, sts_zk); err != nil {
		return nil, err
	}

	sts_zk.Name = resHandle.Name + "-zk"
	sts_zk.Spec.Selector = &metav1.LabelSelector{
		// https://kubernetes.io/docs/concepts/overview/working-with-objects/common-labels/
		MatchLabels: map[string]string{
			"app.kubernetes.io/managed-by": "test-infra",
			TestInfraName:                  resHandle.Name,
			"app.kubernetes.io/component":  "zk",
			"app.kubernetes.io/name":       "kafka-cluster",
		},
	}
	sts_zk.Spec.Template.ObjectMeta.Labels = map[string]string{
		"app.kubernetes.io/managed-by": "test-infra",
		TestInfraName:                  resHandle.Name,
		"app.kubernetes.io/component":  "zk",
		"app.kubernetes.io/name":       "kafka-cluster",
	}

	if resHandle.Spec.ZkSpec.Replicas != 0 {
		sts_zk.Spec.Replicas = &resHandle.Spec.ZkSpec.Replicas
		if sts_zk.Spec.Template.Spec.Containers[0].Env[0].Name != "SERVER_NUM" {
			return nil, errors.New("unexpected yaml template of kafka-zookeeper----env-SERVER_NUM not in the first place")
		}
		sts_zk.Spec.Template.Spec.Containers[0].Env[0] = corev1.EnvVar{
			Name:  "SERVER_NUM",
			Value: fmt.Sprintf("%d", resHandle.Spec.ZkSpec.Replicas),
		}
	}
	if resHandle.Spec.ZkSpec.Image != "" {
		sts_zk.Spec.Template.Spec.Containers[0].Image = resHandle.Spec.ZkSpec.Image
	}
	if len(resHandle.Spec.ZkSpec.Resources.Requests) != 0 {
		for res, request := range resHandle.Spec.ZkSpec.Resources.Requests {
			switch res {
			case apiresource.CPU, apiresource.Memory:
				quantity, err := resource.ParseQuantity(string(request))
				if err != nil {
					return nil, errors.Annotatef(err, "parse quantity on %s", request)
				}
				sts_zk.Spec.Template.Spec.Containers[0].Resources.Requests[corev1.ResourceName(res)] = quantity
				sts_zk.Spec.Template.Spec.Containers[0].Resources.Limits[corev1.ResourceName(res)] = quantity
			case apiresource.Storage:
				quantity, err := resource.ParseQuantity(string(request))
				if err != nil {
					return nil, errors.Annotatef(err, "parse quantity on %s", request)
				}
				sts_zk.Spec.VolumeClaimTemplates[0].Spec.Resources.Requests = corev1.ResourceList{
					corev1.ResourceStorage: quantity,
				}
			}
		}
	}
	return sts_zk, nil
}

func (t KafkaClusterClient) buildSTS_KafkaObj(resHandle *apiresource.KafkaCluster) (*v1.StatefulSet, error) {
	sts_kafka := &v1.StatefulSet{}
	if err := yaml.Unmarshal(kafkaClusterSpec, sts_kafka); err != nil {
		return nil, err
	}

	sts_kafka.Name = resHandle.Name + "-kafka"
	sts_kafka.Spec.Selector = &metav1.LabelSelector{
		// https://kubernetes.io/docs/concepts/overview/working-with-objects/common-labels/
		MatchLabels: map[string]string{
			"app.kubernetes.io/managed-by": "test-infra",
			TestInfraName:                  resHandle.Name,
			"app.kubernetes.io/component":  "kafka",
			"app.kubernetes.io/name":       "kafka-cluster",
		},
	}
	sts_kafka.Spec.Template.ObjectMeta.Labels = map[string]string{
		"app.kubernetes.io/managed-by": "test-infra",
		TestInfraName:                  resHandle.Name,
		"app.kubernetes.io/component":  "kafka",
		"app.kubernetes.io/name":       "kafka-cluster",
	}

	if resHandle.Spec.KafkaSpec.Replicas != 0 {
		sts_kafka.Spec.Replicas = &resHandle.Spec.KafkaSpec.Replicas
	}
	if resHandle.Spec.KafkaSpec.Image != "" {
		sts_kafka.Spec.Template.Spec.Containers[0].Image = resHandle.Spec.KafkaSpec.Image
	}
	if resHandle.Spec.KafkaSpec.ImagePullPolicy != "" {
		sts_kafka.Spec.Template.Spec.Containers[0].ImagePullPolicy = corev1.PullPolicy(resHandle.Spec.KafkaSpec.ImagePullPolicy)
	}
	if resHandle.Spec.KafkaSpec.Command != nil {
		sts_kafka.Spec.Template.Spec.Containers[0].Command = resHandle.Spec.KafkaSpec.Command
	}
	if resHandle.Spec.KafkaSpec.Args != nil {
		sts_kafka.Spec.Template.Spec.Containers[0].Args = resHandle.Spec.KafkaSpec.Args
	}
	if resHandle.Spec.KafkaSpec.NodeSelector != nil {
		sts_kafka.Spec.Template.Spec.NodeSelector = resHandle.Spec.KafkaSpec.NodeSelector
	}
	if len(resHandle.Spec.KafkaSpec.Resources.Requests) != 0 {
		for res, request := range resHandle.Spec.KafkaSpec.Resources.Requests {
			switch res {
			case apiresource.CPU, apiresource.Memory:
				quantity, err := resource.ParseQuantity(string(request))
				if err != nil {
					return nil, errors.Annotatef(err, "parse quantity on %s", request)
				}
				sts_kafka.Spec.Template.Spec.Containers[0].Resources.Requests[corev1.ResourceName(res)] = quantity
				sts_kafka.Spec.Template.Spec.Containers[0].Resources.Limits[corev1.ResourceName(res)] = quantity
			case apiresource.Storage:
				quantity, err := resource.ParseQuantity(string(request))
				if err != nil {
					return nil, errors.Annotatef(err, "parse quantity on %s", request)
				}
				sts_kafka.Spec.VolumeClaimTemplates[0].Spec.Resources.Requests = corev1.ResourceList{
					corev1.ResourceStorage: quantity,
				}
			}
		}
	}

	appendOrOverride := func(env []corev1.EnvVar, key, value string) []corev1.EnvVar {
		item := corev1.EnvVar{Name: key, Value: value}
		for i := range env {
			if env[i].Name == key {
				env[i] = item
				return env
			}
		}
		return append(env, item)
	}

	for k, v := range resHandle.Spec.KafkaSpec.Env {
		sts_kafka.Spec.Template.Spec.Containers[0].Env = appendOrOverride(sts_kafka.Spec.Template.Spec.Containers[0].Env, k, v)
	}

	return sts_kafka, nil
}

func (t KafkaClusterClient) buildKafkaService(resHandleName string, owner *metav1.OwnerReference) (*corev1.Service, *corev1.Service) {
	service_peer := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      resHandleName + "-kafka-peer",
			Namespace: t.Namespace,
			Labels: map[string]string{
				"app.kubernetes.io/managed-by": "test-infra",
				TestInfraName:                  resHandleName,
				"app.kubernetes.io/component":  "kafka",
				"app.kubernetes.io/name":       "kafka-cluster",
			},
			OwnerReferences: []metav1.OwnerReference{*owner},
		},
		// if ClusterIP is `None`, Type must set to `ClusterIP`, else no need to set `ClusterIP`
		Spec: corev1.ServiceSpec{
			Type:      corev1.ServiceTypeClusterIP,
			ClusterIP: corev1.ClusterIPNone,
			Selector: map[string]string{
				"app.kubernetes.io/managed-by": "test-infra",
				TestInfraName:                  resHandleName,
				"app.kubernetes.io/component":  "kafka",
				"app.kubernetes.io/name":       "kafka-cluster",
			},
		},
	}

	service_peer.Spec.Ports = []corev1.ServicePort{
		{
			Name: "server",
			Port: 9092,
		},
	}

	service := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      resHandleName + "-kafka",
			Namespace: t.Namespace,
			Labels: map[string]string{
				"app.kubernetes.io/managed-by": "test-infra",
				TestInfraName:                  resHandleName,
				"app.kubernetes.io/component":  "kafka",
				"app.kubernetes.io/name":       "kafka-cluster",
			},
			OwnerReferences: []metav1.OwnerReference{*owner},
		},
		Spec: corev1.ServiceSpec{
			Type: corev1.ServiceTypeClusterIP,
			Selector: map[string]string{
				"app.kubernetes.io/managed-by": "test-infra",
				TestInfraName:                  resHandleName,
				"app.kubernetes.io/component":  "kafka",
				"app.kubernetes.io/name":       "kafka-cluster",
			},
		},
	}

	service.Spec.Ports = []corev1.ServicePort{
		{
			Name: "server",
			Port: 9092,
		},
	}
	return service_peer, service
}

func (t KafkaClusterClient) buildZKService(resHandleName string, owner *metav1.OwnerReference) (*corev1.Service, *corev1.Service) {
	service_peer := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      resHandleName + "-zk-peer",
			Namespace: t.Namespace,
			Labels: map[string]string{
				"app.kubernetes.io/managed-by": "test-infra",
				TestInfraName:                  resHandleName,
				"app.kubernetes.io/component":  "zk",
				"app.kubernetes.io/name":       "kafka-cluster",
			},
			OwnerReferences: []metav1.OwnerReference{*owner},
		},
		// if ClusterIP is `None`, Type must set to `ClusterIP`, else no need to set `ClusterIP`
		Spec: corev1.ServiceSpec{
			Type:      corev1.ServiceTypeClusterIP,
			ClusterIP: corev1.ClusterIPNone,
			Selector: map[string]string{
				"app.kubernetes.io/managed-by": "test-infra",
				TestInfraName:                  resHandleName,
				"app.kubernetes.io/component":  "zk",
				"app.kubernetes.io/name":       "kafka-cluster",
			},
		},
	}

	service_peer.Spec.Ports = []corev1.ServicePort{
		{
			Name: "server",
			Port: 2888,
		}, {
			Name: "leader-election",
			Port: 3888,
		},
	}

	service := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      resHandleName + "-zk",
			Namespace: t.Namespace,
			Labels: map[string]string{
				"app.kubernetes.io/managed-by": "test-infra",
				TestInfraName:                  resHandleName,
				"app.kubernetes.io/component":  "zk",
				"app.kubernetes.io/name":       "kafka-cluster",
			},
			OwnerReferences: []metav1.OwnerReference{*owner},
		},
		// if ClusterIP is `None`, Type must set to `ClusterIP`, else no need to set `ClusterIP`
		Spec: corev1.ServiceSpec{
			Type: corev1.ServiceTypeClusterIP,
			Selector: map[string]string{
				"app.kubernetes.io/managed-by": "test-infra",
				TestInfraName:                  resHandleName,
				"app.kubernetes.io/component":  "zk",
				"app.kubernetes.io/name":       "kafka-cluster",
			},
		},
	}
	service.Spec.Ports = []corev1.ServicePort{
		{
			Name: "client",
			Port: 2181,
		},
	}

	return service_peer, service
}

type kafkaGetter struct{ handle apiresource.ResourceHandle }

func (t *kafkaGetter) Get() (core.Resource, error) {
	k, err := NewKubeClientFromKubeConfig(t.handle.Testbed, os.Getenv("KUBECONFIG"))
	if err != nil {
		return nil, err
	}
	return NewKafkaClusterClient(k).Get(t.handle)
}

func init() {
	core.RegisterProvider(sdkresource.FQDN(KafkaClusterClient{}),
		func(handle apiresource.ResourceHandle) core.ResourceGetter { return &kafkaGetter{handle: handle} })
}