package framework

import (
	"fmt"
	"time"

	"k8s.io/apimachinery/pkg/util/wait"

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

func (f *Framework) CreateKafkaClusterWithSpecificIMG(testbed *k8s.Testbed) *k8s.KafkaCluster {
	client := k8s.NewKafkaClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name": "kafka-cluster",
  "testbed": "%s",
  "type": "%s",
  "spec": {
    "zkspec": {
      "replicas": 1,
      "image": "k8s.gcr.io/kubernetes-zookeeper:1.0-3.4.10",
      "resources":{
        "requests":{
          "memory":"1Gi",
          "cpu":"500m",
          "storage":"10Gi"
        }
      }
    },
    "kafkaspec": {
      "replicas": 1,
      "imagePullPolicy": "IfNotPresent",
      "image":"registry-mirror.pingcap.net/wurstmeister/kafka:2.13-2.7.1",
      "command":[
        "sh",
        "-c",
        "kafka-server-start.sh /opt/kafka/config/server.properties --override broker.id=${HOSTNAME##*-} --override listeners=PLAINTEXT://:9092 --override zookeeper.connect=$ZK_CONNECT --override log.dir=/var/lib/kafka "
      ],
      "resources":{
        "requests":{
          "memory":"1Gi",
          "cpu":"500m",
          "storage":"1Gi"
        }
      },
      "env": {
        "KAFKA_HEAP_OPTS": "-Xmx512M -Xms512M",
        "KAFKA_OPTS": "-Dlogging.level=INFO"
      }
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeKafkaCluster)))

	ExpectNoError(err)
	return r.(*k8s.KafkaCluster)
}

func (f *Framework) CreateKafkaCluster(testbed *k8s.Testbed) *k8s.KafkaCluster {
	client := k8s.NewKafkaClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name": "kafka-cluster",
  "testbed": "%s",
  "type": "%s",
  "spec": {
    "zkspec": {
      "replicas": 1,
      "resources":{
        "requests":{
          "memory":"1Gi",
          "cpu":"500m",
          "storage":"10Gi"
        }
      }
    },
    "kafkaspec": {
      "replicas": 1,
      "resources":{
        "requests":{
          "memory":"1Gi",
          "cpu":"500m",
          "storage":"1Gi"
        }
      }
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeKafkaCluster)))

	ExpectNoError(err)
	return r.(*k8s.KafkaCluster)
}

func (f *Framework) CreateKafkaClusterWithMultiBrokers(testbed *k8s.Testbed, zkNum, kafkaNum int) *k8s.KafkaCluster {
	client := k8s.NewKafkaClusterClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name": "kafka-cluster",
  "testbed": "%s",
  "type": "%s",
  "spec": {
    "zkspec": {
      "replicas": %d,
      "resources":{
        "requests":{
          "memory":"1Gi",
          "cpu":"500m",
          "storage":"10Gi"
        }
      }
    },
    "kafkaspec": {
      "replicas": %d,
      "resources":{
        "requests":{
          "memory":"1Gi",
          "cpu":"500m",
          "storage":"1Gi"
        }
      }
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeKafkaCluster, zkNum, kafkaNum)))

	ExpectNoError(err)
	return r.(*k8s.KafkaCluster)
}

func (f *Framework) CreateKafkacatWorkloadNode(testbed *k8s.Testbed) *k8s.WorkloadNode {
	client := k8s.NewWorkloadNodeClient(f.KubeClientSetPack)

	r, err := client.Create([]byte(fmt.Sprintf(`{
  "name":"kafkacat",
  "testbed":"%s",
  "type":"%s",
  "spec":{
    "container":{
      "name":"kafkacat",
      "imagePullPolicy": "IfNotPresent",
      "image":"registry-mirror.pingcap.net/edenhill/kcat:1.7.0",
      "command":[
        "tail",
        "-f",
        "/dev/null"
      ],
      "resources":{
        "requests":{
          "memory":"300Mi",
          "cpu":"250m"
        }
      }
    }
  }
}
`, testbed.Name, apiresource.ResourceTypeWorkloadNode)))

	ExpectNoError(err)
	return r.(*k8s.WorkloadNode)
}

func (f *Framework) GetKafkaCluster(handle apiresource.ResourceHandle) *k8s.KafkaCluster {
	client := k8s.NewKafkaClusterClient(f.KubeClientSetPack)
	r, err := client.Get(handle)
	ExpectNoError(err)
	return r.(*k8s.KafkaCluster)
}

func (f *Framework) DeleteKafkaClusterWaited(kafka *k8s.KafkaCluster) {
	client := k8s.NewKafkaClusterClient(f.KubeClientSetPack)
	ExpectNoError(client.Delete(kafka.ResourceHandle))

	err := wait.PollImmediate(poll, f.DeletionTimeout, func() (done bool, err error) {
		_, err = client.Get(kafka.ResourceHandle)
		if resource.IsResourceNotFoundError(err) {
			return true, nil
		}
		return false, nil
	})
	ExpectNoError(err)
}

func (f *Framework) WaitForKafkaClusterReady(kafka *k8s.KafkaCluster, timeout time.Duration) error {
	if err := wait.PollImmediate(poll, timeout, func() (done bool, err error) {
		if err := kafka.Sync(); err != nil {
			return false, err
		}
		return kafka.Ready(), nil
	}); err != nil {
		return fmt.Errorf("wait for kafka cluster ready occur an error: %v", err)
	}
	return nil
}