{{- $advertisedListenersOverride := first (pluck "advertised.listeners" .Values.configurationOverrides) }}
apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: {{ include "kafka.fullname" . }}
  labels:
  {{- include "kafka.broker.labels" . | nindent 4 }}
spec:
  selector:
    matchLabels:
    {{- include "kafka.broker.matchLabels" . | nindent 6 }}
  serviceName: {{ include "kafka.fullname" . }}-headless
  podManagementPolicy: {{ .Values.podManagementPolicy }}
  updateStrategy:
{{ toYaml .Values.updateStrategy | indent 4 }}
  replicas: {{ default 3 .Values.replicas }}
  template:
    metadata:
{{- if or .Values.podAnnotations (and .Values.prometheus.jmx.enabled (not .Values.prometheus.operator.enabled)) }}
      annotations:
{{- if and .Values.prometheus.jmx.enabled  (not .Values.prometheus.operator.enabled) }}
        prometheus.io/scrape: "true"
        prometheus.io/port: {{ .Values.prometheus.jmx.port | quote }}
{{- end }}
{{- if .Values.podAnnotations }}
{{ toYaml .Values.podAnnotations | indent 8 }}
{{- end }}
{{- end }}
      labels:
      {{- include "kafka.broker.labels" . | nindent 8 }}
        {{- if .Values.podLabels }}
        ## Custom pod labels
{{ toYaml .Values.podLabels | indent 8 }}
        {{- end }}
    spec:
{{- if .Values.schedulerName }}
      schedulerName: "{{ .Values.schedulerName }}"
{{- end }}
{{- if .Values.serviceAccountName }}
      serviceAccountName: {{ .Values.serviceAccountName }}
{{- end }}
{{- if .Values.priorityClassName }}
      priorityClassName: "{{ .Values.priorityClassName }}"
{{- end }}
{{- if .Values.tolerations }}
      tolerations:
{{ toYaml .Values.tolerations | indent 8 }}
{{- end }}
{{- if .Values.affinity }}
      affinity:
{{ toYaml .Values.affinity | indent 8 }}
{{- end }}
{{- if .Values.nodeSelector }}
      nodeSelector:
{{ toYaml .Values.nodeSelector | indent 8 }}
{{- end }}
      containers:
      {{- if .Values.prometheus.jmx.enabled }}
      - name: metrics
        image: "{{ .Values.prometheus.jmx.image }}:{{ .Values.prometheus.jmx.imageTag }}"
        command:
        - sh
        - -exc
        - |
          trap "exit 0" TERM; \
          while :; do \
          java \
          -XX:+UnlockExperimentalVMOptions \
          -XX:+UseCGroupMemoryLimitForHeap \
          -XX:MaxRAMFraction=1 \
          -XshowSettings:vm \
          -jar \
          jmx_prometheus_httpserver.jar \
          {{ .Values.prometheus.jmx.port | quote }} \
          /etc/jmx-kafka/jmx-kafka-prometheus.yml & \
          wait $! || sleep 3; \
          done
        ports:
        - containerPort: {{ .Values.prometheus.jmx.port }}
          name: prometheus
        resources:
{{ toYaml .Values.prometheus.jmx.resources | indent 10 }}
        volumeMounts:
        - name: jmx-config
          mountPath: /etc/jmx-kafka
      {{- end }}
      - name: {{ include "kafka.name" . }}-broker
        image: "{{ .Values.image }}:{{ .Values.imageTag }}"
        imagePullPolicy: "{{ .Values.imagePullPolicy }}"
        livenessProbe:
          exec:
            command:
              - sh
              - -ec
              - /usr/bin/jps | /bin/grep -q SupportedKafka
          {{- if not .Values.livenessProbe }}
          initialDelaySeconds: 30
          timeoutSeconds: 5
          {{- else }}
          initialDelaySeconds: {{ .Values.livenessProbe.initialDelaySeconds | default 30}}
          {{- if .Values.livenessProbe.periodSeconds }}
          periodSeconds: {{ .Values.livenessProbe.periodSeconds }}
          {{- end }}
          timeoutSeconds: {{ .Values.livenessProbe.timeoutSeconds | default 5}}
          {{- if .Values.livenessProbe.successThreshold }}
          successThreshold: {{ .Values.livenessProbe.successThreshold }}
          {{- end }}
          {{- if .Values.livenessProbe.failureThreshold }}
          failureThreshold: {{ .Values.livenessProbe.failureThreshold }}
          {{- end }}
          {{- end }}
        readinessProbe:
          tcpSocket:
            port: kafka
          initialDelaySeconds: {{ .Values.readinessProbe.initialDelaySeconds }}
          periodSeconds: {{ .Values.readinessProbe.periodSeconds }}
          timeoutSeconds: {{ .Values.readinessProbe.timeoutSeconds }}
          successThreshold: {{ .Values.readinessProbe.successThreshold }}
          failureThreshold: {{ .Values.readinessProbe.failureThreshold }}
        ports:
        - containerPort: 9092
          name: kafka
        {{- if .Values.external.enabled }}
          {{- $replicas := .Values.replicas | int }}
          {{- $root := . }}
          {{- range $i, $e := until $replicas }}
        - containerPort: {{ add $root.Values.external.firstListenerPort $i }}
          name: external-{{ $i }}
          {{- end }}
        {{- end }}
        {{- if .Values.prometheus.jmx.enabled }}
        - containerPort: {{ .Values.jmx.port }}
          name: jmx
        {{- end }}
        {{- if .Values.additionalPorts }}
{{ toYaml .Values.additionalPorts | indent 8 }}
        {{- end }}
        resources:
{{ toYaml .Values.resources | indent 10 }}
        env:
        {{- if .Values.prometheus.jmx.enabled }}
        - name: JMX_PORT
          value: "{{ .Values.jmx.port }}"
        {{- end }}
        - name: POD_IP
          valueFrom:
            fieldRef:
              fieldPath: status.podIP
        - name: POD_NAME
          valueFrom:
            fieldRef:
              fieldPath: metadata.name
        - name: POD_NAMESPACE
          valueFrom:
            fieldRef:
              fieldPath: metadata.namespace
        - name: KAFKA_HEAP_OPTS
          value: {{ .Values.kafkaHeapOptions }}
        - name: KAFKA_OFFSETS_TOPIC_REPLICATION_FACTOR
          value: {{ include "kafka.replication.factor" . | quote }}
        {{- if not (hasKey .Values.configurationOverrides "zookeeper.connect") }}
        - name: KAFKA_ZOOKEEPER_CONNECT
          value: {{ include "zookeeper.url" . | quote }}
        {{- end }}
        {{- if not (hasKey .Values.configurationOverrides "log.dirs") }}
        - name: KAFKA_LOG_DIRS
          value: {{ printf "%s/%s" .Values.persistence.mountPath .Values.logSubPath | quote }}
        {{- end }}
        {{- range $key, $value := .Values.configurationOverrides }}
        - name: {{ printf "KAFKA_%s" $key | replace "." "_" | upper | quote }}
          value: {{ $value | quote }}
        {{- end }}
        {{- if .Values.jmx.port }}
        - name: KAFKA_JMX_PORT
          value: "{{ .Values.jmx.port }}"
        {{- end }}
        {{- range $secret := .Values.secrets }}
          {{- if not $secret.mountPath }}
            {{- range $key := $secret.keys }}
        - name: {{ (print ($secret.name | replace "-" "_") "_" $key) | upper }}
          valueFrom:
            secretKeyRef:
              name: {{ $secret.name }}
              key: {{ $key }}
            {{- end }}
          {{- end }}
        {{- end }}
        {{- range $key, $value := .Values.envOverrides }}
        - name: {{ printf "%s" $key | replace "." "_" | upper | quote }}
          value: {{ $value | quote }}
        {{- end }}
        # This is required because the Downward API does not yet support identification of
        # pod numbering in statefulsets. Thus, we are required to specify a command which
        # allows us to extract the pod ID for usage as the Kafka Broker ID.
        # See: https://github.com/kubernetes/kubernetes/issues/31218
        command:
        - sh
        - -exc
        - |
          unset KAFKA_PORT && \
          export KAFKA_BROKER_ID=${POD_NAME##*-} && \
          {{- if eq .Values.external.type "LoadBalancer" }}
          export LOAD_BALANCER_IP=$(echo '{{ .Values.external.loadBalancerIP }}' | tr -d '[]' | cut -d ' ' -f "$(($KAFKA_BROKER_ID + 1))") && \
          {{- end }}
          {{- if eq .Values.external.type "NodePort" }}
          export KAFKA_ADVERTISED_LISTENERS=PLAINTEXT://${POD_IP}:9092{{ if kindIs "string" $advertisedListenersOverride }}{{ printf ",%s" $advertisedListenersOverride }}{{ end }} && \
          {{- else }}
          export KAFKA_ADVERTISED_LISTENERS=PLAINTEXT://${POD_NAME}.{{ include "kafka.fullname" . }}-headless.${POD_NAMESPACE}.svc.cluster.local:9092{{ if kindIs "string" $advertisedListenersOverride }}{{ printf ",%s" $advertisedListenersOverride }}{{ end }} && \
          {{- end }}
          exec /etc/confluent/docker/run
        volumeMounts:
        - name: datadir
          mountPath: {{ .Values.persistence.mountPath | quote }}
        {{- range $secret := .Values.secrets }}
          {{- if $secret.mountPath }}
            {{- if $secret.keys }}
              {{- range $key := $secret.keys }}
        - name: {{ include "kafka.fullname" $ }}-{{ $secret.name }}
          mountPath: {{ $secret.mountPath }}/{{ $key }}
          subPath: {{ $key }}
          readOnly: true
              {{- end }}
            {{- else }}
        - name: {{ include "kafka.fullname" $ }}-{{ $secret.name }}
          mountPath: {{ $secret.mountPath }}
          readOnly: true
            {{- end }}
          {{- end }}
        {{- end }}
      volumes:
      {{- if not .Values.persistence.enabled }}
      - name: datadir
        emptyDir: {}
      {{- end }}
      {{- if .Values.prometheus.jmx.enabled }}
      - name: jmx-config
        configMap:
          {{- if .Values.jmx.configMap.overrideName }}
          name: {{ .Values.jmx.configMap.overrideName }}
          {{- else }}
          name: {{ include "kafka.fullname" . }}-metrics
          {{- end }}
      {{- end }}
      {{- if .Values.securityContext }}
      securityContext:
{{ toYaml .Values.securityContext | indent 8 }}
      {{- end }}
      {{- range .Values.secrets }}
      {{- if .mountPath }}
      - name: {{ include "kafka.fullname" $ }}-{{ .name }}
        secret:
          secretName: {{ .name }}
      {{- end }}
      {{- end }}
      terminationGracePeriodSeconds: {{ .Values.terminationGracePeriodSeconds }}
  {{- if .Values.persistence.enabled }}
  volumeClaimTemplates:
  - metadata:
      name: datadir
    spec:
      accessModes: [ "ReadWriteOnce" ]
      resources:
        requests:
          storage: {{ .Values.persistence.size }}
      {{- if .Values.persistence.storageClass }}
      {{- if (eq "-" .Values.persistence.storageClass) }}
      storageClassName: ""
      {{- else }}
      storageClassName: "{{ .Values.persistence.storageClass }}"
      {{- end }}
      {{- end }}
  {{- end }}
