---
apiVersion: apps/v1
kind: StatefulSet
metadata:
  name: {{ .Chart.Name }}
  labels:
    app: {{ .Chart.Name }}
spec:
  serviceName: {{ .Chart.Name }}
  minReadySeconds: 10
  updateStrategy:
    type: RollingUpdate
  replicas: 1
  selector:
    matchLabels:
      app: {{ .Chart.Name }}
  template:
    metadata:
      labels:
        app: {{ .Chart.Name }}
    spec:
      affinity:
        podAntiAffinity:
          preferredDuringSchedulingIgnoredDuringExecution:
          - podAffinityTerm:
              labelSelector:
                matchLabels:
                  app: {{ .Chart.Name }}
              topologyKey: kubernetes.io/hostname
            weight: 100
      {{- if and .Values.gpu.enabled .Values.gpu.nodeSelector }}
      nodeSelector: {{ toYaml .Values.gpu.nodeSelector | nindent 8 }}
      {{- end }}
      {{- if and .Values.gpu.enabled .Values.gpu.tolerations }}
      tolerations: {{ toYaml .Values.gpu.tolerations | nindent 6 }}
      {{- end }}
      {{- with .Values.imagePullSecrets }}
      imagePullSecrets:
        {{- toYaml . | nindent 8 }}
      {{- end }}
      containers:
        - name: {{ .Chart.Name }}
          image: {{ .Values.image.ollama }}
          imagePullPolicy: IfNotPresent
          volumeMounts:
            - name: data
              mountPath: /root/.ollama
          ports:
            - name: http
              containerPort: {{ .Values.port }}
              protocol: TCP
          env:
            - name: OLLAMA_HOST
              value: "0.0.0.0:{{ .Values.port }}"
          {{- range $n, $v := .Values.envVars }}
            - name: {{ $n }}
              value: {{ $v | quote }}
          {{- end }}
          livenessProbe:
            httpGet:
              path: /
              port: http
            initialDelaySeconds: 15
            periodSeconds: 9
            timeoutSeconds: 6
            failureThreshold: 3
          readinessProbe:
            httpGet:
              path: /
              port: http
            initialDelaySeconds: 15
            periodSeconds: 5
            timeoutSeconds: 6
            failureThreshold: 3
          resources: {{ dig .Values.env .Values.resources._default .Values.resources | toYaml | nindent 12 }}
          {{- if or .Values.models.pull .Values.models.run }}
          lifecycle:
            postStart:
              exec:
                command:
                  - /bin/sh
                  - -c
                  - |
                    while ! /bin/ollama ps > /dev/null 2>&1; do
                      sleep 5
                    done

                    {{- $allModels := list -}}

                    {{- if .Values.models.pull }}
                      {{- range .Values.models.pull }}

                        {{- if contains ":" . }}
                          {{- $allModels = append $allModels . }}
                        {{- else }}
                          {{- $allModels = append $allModels (printf "%s:latest" .) }}
                        {{- end }}

                        /bin/ollama pull {{ternary "--insecure" "" $.Values.insecure | toString }} {{ . }}
                        {{- end }}
                      {{- end }}

                      {{- if .Values.models.run }}
                        {{- range .Values.models.run }}

                          {{- if contains ":" . }}
                            {{- $allModels = append $allModels . }}
                          {{- else }}
                            {{- $allModels = append $allModels (printf "%s:latest" .) }}
                          {{- end }}

                          /bin/ollama run {{ . }}
                        {{- end }}
                      {{- end }}

                      {{- if .Values.models.clean }}
                        /bin/ollama list | awk 'NR>1 {print $1}' | while read model; do
                          echo "{{ $allModels | join " " }}" | tr ' ' '\n' | grep -Fqx "$model" || /bin/ollama rm "$model"
                        done
                      {{- end }}
          {{- end }}
  volumeClaimTemplates:
  - metadata:
      name: data
    spec:
      accessModes: [ "ReadWriteOnce" ]
      storageClassName: {{ dig .Values.env .Values.storage.class._default .Values.storage.class | quote }}
      resources:
        requests:
          storage: {{ dig .Values.env .Values.storage.size._default .Values.storage.size | quote }}
---
apiVersion: v1
kind: Service
metadata:
  name: {{ .Chart.Name }}
  labels:
    app: {{ .Chart.Name }}
spec:
  type: ClusterIP
  ports:
    - port: {{ .Values.port }}
      targetPort: http
      protocol: TCP
      name: http
  selector:
    app: {{ .Chart.Name }}