infra/modules/nixos/services/ollama_init_custom_models.nix

{ pkgs, ... }: {
  systemd.services.init-ollama-model = {
    description = "Initialize LLM models with extra context in Ollama Docker";
    
    # On s'assure que Docker tourne avant de lancer ce script
    after = [ "docker.service" ];
    wantedBy = [ "multi-user.target" ];
    
    script = ''
      # Fonction de création asynchrone pour ne pas bloquer le démarrage
      (
        echo "Starting asynchronous Ollama initialization..."
        
        # Attente d'Ollama (maximum 120 secondes pour éviter une boucle infinie)
        TIMEOUT=60
        COUNT=0
        while ! ${pkgs.curl}/bin/curl -s -f http://127.0.0.1:11434/api/tags > /dev/null; do
          if [ $COUNT -ge $TIMEOUT ]; then
            echo "Ollama did not become ready in time. Exiting."
            exit 1
          fi
          echo "Waiting for Ollama API to be reachable..."
          sleep 5
          COUNT=$((COUNT + 5))
        done

        create_model_if_missing() {
          local model_name=$1
          local base_model=$2
          
          # Vérification robuste via l'API HTTP d'Ollama plutôt que docker exec (évite les conflits de tty)
          if ! ${pkgs.curl}/bin/curl -s http://127.0.0.1:11434/api/tags | ${pkgs.jq}/bin/jq -e ".models[] | select(.name == \"$model_name\")" > /dev/null; then
            echo "$model_name not found, creating from $base_model..."
            
            # Utilisation d'un fichier temporaire sur l'hôte pour l'injecter proprement dans Docker
            TMP_FILE=$(mktemp)
            cat <<EOF > "$TMP_FILE"
FROM $base_model
TEMPLATE """{{- if .System }}
[SYSTEM_PROMPT]
{{ .System }}
[/SYSTEM_PROMPT]
{{- end }}
{{- range .Messages }}
{{- if eq .Role "user" }}
[INST]
{{ .Content }}
[/INST]
{{- else if eq .Role "assistant" }}
{{ .Content }}
{{- end }}
{{- end }}"""
PARAMETER num_ctx 131072
PARAMETER num_predict 4096
PARAMETER num_keep 1024
PARAMETER repeat_penalty 1.1
PARAMETER top_k 40
PARAMETER stop "[INST]"
PARAMETER stop "[/INST]"
PARAMETER stop "</s>"
EOF

            # Copie et création dans le conteneur
            ${pkgs.docker}/bin/docker cp "$TMP_FILE" ollama:/tmp/model.modelfile
            ${pkgs.docker}/bin/docker exec ollama ollama create "$model_name" -f /tmp/model.modelfile
            ${pkgs.docker}/bin/docker exec ollama rm /tmp/model.modelfile
            rm -f "$TMP_FILE"
          else
            echo "$model_name already exists, skipping."
          fi
        }

        # Create Nemotron
        create_model_if_missing "nemotron-3-nano:30b-128k" "nemotron-3-nano:30b"
        
        # Create Devstral
        create_model_if_missing "devstral-small-2:24b-128k" "devstral-small-2:24b" 
        
      ) &
    '';

    serviceConfig = {
      Type = "forking"; # Permet à systemd de savoir que le script passe en arrière-plan via '&'
      User = "root";
    };
  };
}