package k8s

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"math/rand"
	"net/url"

	"github.com/hashicorp/go-multierror"
	"github.com/pingcap/errors"
	"go.uber.org/zap"
	v1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/labels"

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

var (
	_ resource.KafkaCluster = &KafkaCluster{}
)

type KafkaCluster struct {
	KubeClientSet
	// api object
	apiresource.KafkaCluster
	// k8s object
	sts_zk             *v1.StatefulSet
	sts_kafka          *v1.StatefulSet
	service_zk         *corev1.Service
	service_zk_peer    *corev1.Service
	service_kafka      *corev1.Service
	service_kafka_peer *corev1.Service
}

func (t *KafkaCluster) Handle() apiresource.ResourceHandle {
	return t.KafkaCluster.ResourceHandle
}

func (t *KafkaCluster) SinkEndpoint() (*url.URL, error) {
	return t.ServiceURL()
}

// ServiceURL can be used to get kafka endpoint
func (t *KafkaCluster) ServiceURL() (*url.URL, error) {
	// svc-name.ns-name.svc
	return url.Parse(fmt.Sprintf("kafka://%s-kafka.%s:9092", t.Name, t.Namespace))
}

// ZookeeperURL can be used to get zookeeper endpoint
func (t *KafkaCluster) ZookeeperURL() (*url.URL, error) {
	// svc-name.ns-name.svc:2181
	return url.Parse(fmt.Sprintf("%s-zk.%s.svc:2181", t.Name, t.Namespace))
}

func (t *KafkaCluster) PeerURLs() ([]*url.URL, error) {
	var (
		format   = "%s-kafka-%d.%s-kafka-peer.%s:9092"
		replicas = 1
		urls     = make([]*url.URL, replicas)
		err      error
	)

	for i := 0; i < replicas; i++ {
		urls[i], err = url.Parse(fmt.Sprintf(format, t.Name, i, t.Name, t.Namespace))
		if err != nil {
			return nil, err
		}
	}
	return urls, nil
}

func (t *KafkaCluster) Exec(opts resource.KafkaClusterExecOptions) (stdout string, stderr string, exitCode int, err error) {
	podName := fmt.Sprintf("%s-%s-%d", t.Name, opts.Target.String(), opts.Peer)
	stdout, stderr, ec, err := t.ExecCommandInContainerWithFullOutput(podName, opts.Target.String(), "/bin/sh", "-c", opts.Command)
	if err != nil {
		sugar.Error("Exec error", zap.Error(err))
		err = errors.Trace(err)
	}
	if ec != nil {
		exitCode = *ec
	} else {
		sugar.Warnw("cannot get the exit code, return -1")
		exitCode = -1
	}
	return
}

type brokerInfo struct {
	Version   int    `json:"version"`
	ID        int    `json:"brokerid"`
	Timestamp string `json:"timestamp"`
}

// Unmarshal the data to initialize the broker info, data would have the format
// WatchedEvent state:SyncConnected type:None path:null
// {"version":1,"brokerid":0,"timestamp":"1646289830437"}
func (b *brokerInfo) Unmarshal(data []byte) error {
	index := bytes.Index(data, []byte("WatchedEvent"))
	if index == -1 {
		sugar.Warnw("unexpected data for broker info", zap.String("data", string(data)))
		return fmt.Errorf("cannot find WatchedEvent when try to fetch kafka broker info")
	}
	data = data[index:]
	lines := bytes.Split(data, []byte("\n"))
	if len(lines) != 2 {
		sugar.Warnw("unexpected data for broker info", zap.String("data", string(data)))
		return fmt.Errorf("unexpected information when try to fetch kafka broker info")
	}

	data = lines[1]
	if err := json.Unmarshal(data, b); err != nil {
		sugar.Warnw("unmarshal broker info failed", zap.Error(err), zap.String("data", string(data)))
		return err
	}
	return nil
}

func (t *KafkaCluster) GetController() (*brokerInfo, error) {
	zkURL, err := t.ZookeeperURL()
	if err != nil {
		return nil, err
	}
	cmd := fmt.Sprintf("zookeeper-shell.sh %s get /controller", zkURL.String())
	option := resource.KafkaClusterExecOptions{
		Command: cmd,
		Target:  resource.Kafka,
		Peer:    0,
	}

	stdout, stderr, exitCode, err := t.Exec(option)
	if err != nil {
		return nil, err
	}
	if exitCode == -1 {
		return nil, fmt.Errorf("get kafka controller failed, ExitCode=%+v", exitCode)
	}
	if stderr != "" {
		sugar.Warnw("stderr is not empty", zap.String("stderr", stderr))
	}

	result := new(brokerInfo)
	if err := result.Unmarshal([]byte(stdout)); err != nil {
		return nil, err
	}

	return result, nil
}

func (t *KafkaCluster) Ready() bool {
	return t.KafkaCluster.Status.Phase == apiresource.KafkaClusterReady
}

// Sync syncs from the k8s object to api object
func (t *KafkaCluster) Sync() error {
	var errs error

	sts_zk, err := t.ClientSet.AppsV1().StatefulSets(t.Namespace).Get(context.TODO(), t.ResourceHandle.Name+"-zk", metav1.GetOptions{})
	if err != nil {
		errs = multierror.Append(errs, err)
		sts_zk = nil
	}
	sts_kafka, err := t.ClientSet.AppsV1().StatefulSets(t.Namespace).Get(context.TODO(), t.ResourceHandle.Name+"-kafka", metav1.GetOptions{})
	if err != nil {
		errs = multierror.Append(errs, err)
		sts_kafka = nil
	}
	service_zk, err := t.ClientSet.CoreV1().Services(t.Namespace).Get(context.TODO(), t.ResourceHandle.Name+"-zk", metav1.GetOptions{})
	if err != nil {
		errs = multierror.Append(errs, err)
		service_zk = nil
	}
	service_zk_peer, err := t.ClientSet.CoreV1().Services(t.Namespace).Get(context.TODO(), t.ResourceHandle.Name+"-zk-peer", metav1.GetOptions{})
	if err != nil {
		errs = multierror.Append(errs, err)
		service_zk_peer = nil
	}
	service_kafka, err := t.ClientSet.CoreV1().Services(t.Namespace).Get(context.TODO(), t.ResourceHandle.Name+"-kafka", metav1.GetOptions{})
	if err != nil {
		errs = multierror.Append(errs, err)
		service_kafka = nil
	}
	service_kafka_peer, err := t.ClientSet.CoreV1().Services(t.Namespace).Get(context.TODO(), t.ResourceHandle.Name+"-kafka-peer", metav1.GetOptions{})
	if err != nil {
		errs = multierror.Append(errs, err)
		service_kafka_peer = nil
	}
	pvcs, err := t.ClientSet.CoreV1().PersistentVolumeClaims(t.Namespace).List(context.TODO(), metav1.ListOptions{
		LabelSelector: labels.Set{TestInfraName: t.Name}.AsSelector().String(),
	})
	if err != nil {
		return multierror.Append(errs, err)
	}
	if len(pvcs.Items) == 0 && sts_zk == nil && sts_kafka == nil && service_zk == nil && service_zk_peer == nil && service_kafka_peer == nil {
		return resource.ErrorResourceNotFound{ResourceHandle: t.Handle()}
	}
	if errs != nil {
		return errs
	}

	t.sts_zk = sts_zk
	t.sts_kafka = sts_kafka
	t.service_zk = service_zk
	t.service_zk_peer = service_zk_peer
	t.service_kafka = service_kafka
	t.service_kafka_peer = service_kafka_peer

	t.KafkaCluster.Spec = apiresource.KafkaClusterSpec{
		ZkSpec: apiresource.ZooKeeperComp{
			Replicas: *sts_zk.Spec.Replicas,
			Image:    sts_zk.Spec.Template.Spec.Containers[0].Image,
			Resources: apiresource.ResourceRequirements{
				Requests: make(map[apiresource.ResourceName]apiresource.Quantity),
			},
		},
		KafkaSpec: apiresource.KafkaComp{
			Replicas:        *sts_kafka.Spec.Replicas,
			Image:           sts_kafka.Spec.Template.Spec.Containers[0].Image,
			ImagePullPolicy: apiresource.ImagePullPolicy(sts_kafka.Spec.Template.Spec.Containers[0].ImagePullPolicy),

			Command: sts_kafka.Spec.Template.Spec.Containers[0].Command,
			Args:    sts_kafka.Spec.Template.Spec.Containers[0].Args,
			Resources: apiresource.ResourceRequirements{
				Requests: make(map[apiresource.ResourceName]apiresource.Quantity),
			},
		},
	}
	for res, quantity := range sts_zk.Spec.Template.Spec.Containers[0].Resources.Requests {
		t.KafkaCluster.Spec.ZkSpec.Resources.Requests[apiresource.ResourceName(res)] = apiresource.Quantity(quantity.String())
	}
	for res, quantity := range sts_kafka.Spec.Template.Spec.Containers[0].Resources.Requests {
		t.KafkaCluster.Spec.KafkaSpec.Resources.Requests[apiresource.ResourceName(res)] = apiresource.Quantity(quantity.String())
	}
	if len(sts_zk.Spec.VolumeClaimTemplates) != 0 {
		for _, volume := range sts_zk.Spec.VolumeClaimTemplates {
			if volume.Name == "data-zk" {
				if quantity, ok := volume.Spec.Resources.Requests[corev1.ResourceStorage]; ok {
					t.KafkaCluster.Spec.ZkSpec.Resources.Requests[apiresource.Storage] = apiresource.Quantity(quantity.String())
				}
			}
		}
	}
	if len(sts_kafka.Spec.VolumeClaimTemplates) != 0 {
		for _, volume := range sts_kafka.Spec.VolumeClaimTemplates {
			if volume.Name == "data-kafka" {
				if quantity, ok := volume.Spec.Resources.Requests[corev1.ResourceStorage]; ok {
					t.KafkaCluster.Spec.ZkSpec.Resources.Requests[apiresource.Storage] = apiresource.Quantity(quantity.String())
				}
			}
		}
	}
	if sts_kafka.Status.ReadyReplicas == t.Spec.KafkaSpec.Replicas {
		t.KafkaCluster.Status.Phase = apiresource.KafkaClusterReady
		t.KafkaCluster.ResourceHandle.Ready = &[]bool{true}[0]
	} else if sts_zk.Status.ReadyReplicas == t.Spec.ZkSpec.Replicas {
		t.KafkaCluster.Status.Phase = apiresource.KafkaClusterNotReady
		t.KafkaCluster.ResourceHandle.Ready = &[]bool{false}[0]
	} else {
		t.KafkaCluster.Status.Phase = apiresource.KafkaClusterZKNotReady
		t.KafkaCluster.ResourceHandle.Ready = &[]bool{false}[0]
	}

	return nil
}

func (t *KafkaCluster) ChaosClient() resource.ChaosClient {
	return &ChaosClient{t.KubeClientSet}
}

func (t *KafkaCluster) ChaosSelect(selector resource.KafkaClusterSelectorOpts) ([]resource.ChaosSelector, error) {
	var selectors []resource.ChaosSelector
	if selector.Controller {
		controller, err := t.GetController()
		if err != nil {
			return nil, err
		}
		return t.selectByID(controller.ID)
	}

	if selector.RandomSameOne {
		pods, err := t.KubeClientSet.ClientSet.CoreV1().Pods(t.Namespace).List(context.Background(), metav1.ListOptions{
			LabelSelector: labels.Set{"app.kubernetes.io/component": "kafka", TestInfraName: t.Name}.AsSelector().String(),
		})
		if err != nil {
			return nil, err
		}
		if len(pods.Items) == 0 {
			return nil, fmt.Errorf("kafka pod not found")
		}
		target := rand.Intn(len(pods.Items))
		selectors = append(selectors, &PodSelector{Pod: &pods.Items[target]})
		return selectors, nil
	}

	if selector.All {
		pods, err := t.KubeClientSet.ClientSet.CoreV1().Pods(t.Namespace).List(context.Background(), metav1.ListOptions{
			LabelSelector: labels.Set{"app.kubernetes.io/component": "kafka", TestInfraName: t.Name}.AsSelector().String(),
		})
		if err != nil {
			return nil, err
		}
		for i := range pods.Items {
			selectors = append(selectors, &PodSelector{Pod: &pods.Items[i]})
		}

		return selectors, nil
	}

	return t.selectByID(selector.Peer)
}

func (t *KafkaCluster) selectByID(id int) ([]resource.ChaosSelector, error) {
	pod, err := t.KubeClientSet.ClientSet.CoreV1().Pods(t.Namespace).Get(context.Background(), fmt.Sprintf("%s-kafka-%d", t.Name, id), metav1.GetOptions{})
	if err != nil {
		return nil, err
	}
	return []resource.ChaosSelector{&PodSelector{pod}}, nil
}