#!/bin/zsh
# Coded by Korben from korben.info

# install_llm_autosuggest.sh
# This script installs the llm-mlx plugin for LLM and sets up a zsh widget that provides autosuggestions using the local LLM model.

# Check if the 'llm' command is available
if ! command -v llm &> /dev/null; then
    echo "Error: 'llm' command not found. Please install LLM using one of the following methods:"
    echo "  - uv tool install llm
  - pipx install llm
  - brew install llm
  - pip install llm"
    exit 1
fi

# Install the llm-mlx plugin (macOS only)
echo "Installing llm-mlx plugin..."
llm install llm-mlx

# Download and register the chosen model (Llama 3.2 3B Instruct 4bit)
echo "Downloading model..."
llm mlx download-model mlx-community/Qwen2.5-Coder-1.5B-Instruct-Q6

# Optionally, you can set an alias for the model if desired
# For example: llm aliases set qwen mlx-community/Llama-3.2-3B-Instruct-4bit
# Uncomment the following line to set the alias:
llm aliases set qwen mlx-community/Qwen2.5-Coder-1.5B-Instruct-Q6

# Append the LLM MLX based autosuggestion widget to ~/.zshrc
# This widget binds Control+Z to transform the current command line using the installed LLM model

echo "Updating LLM-based autosuggestion widget in ~/.zshrc..."

# Remove existing configuration if present
sed -i '' '/# === LLM MLX Autosuggestion Widget Start ===/,/# === LLM MLX Autosuggestion Widget End ===/d' ~/.zshrc

# Add new configuration
cat << 'EOF' >> ~/.zshrc

# === LLM MLX Autosuggestion Widget Start ===
_llm_autosuggest() {
  if [[ -n "$BUFFER" ]]; then
    local prev_cmd="$BUFFER"
    # Escape the command properly for shell
    local escaped_cmd=$(printf %q "$prev_cmd")
    # Show progress indicator
    BUFFER+=$' ⌛'
    zle -I && zle redisplay
    # Use the model alias 'qwen' if set, otherwise default to full model identifier
    local suggestion
    if command -v llm &> /dev/null; then
      raw_suggestion=$(llm -m qwen <<< "Fournis uniquement la ligne de commande demandée, sans aucun texte explicatif ou commentaire. Ne mets pas la ligne de commande entre quote. Système: macOS. Contexte dossier: $PWD. Vérifie que la commande est valide sous macOS. Commande: $escaped_cmd")
      # Fallback if alias 'qwen' didn't work, use full identifier
      if [ -z "$raw_suggestion" ]; then
        raw_suggestion=$(llm -m mlx-community/Qwen2.5-Coder-1.5B-Instruct-Q6 <<< "Fournis uniquement la ligne de commande demandée, sans aucun texte explicatif ou commentaire. Ne mets pas la ligne de commande entre quote. Système: macOS. Contexte dossier: $PWD. Vérifie que la commande est valide sous macOS. Commande: $escaped_cmd")
      fi
      # Remove content between <think> tags using sed
      suggestion=$(echo "$raw_suggestion" | sed '/<think>/,/<\/think>/d')
    else
      suggestion="$prev_cmd"
    fi
    # Extract only the first line of the suggestion
    # Get first non-empty line after XML filtering
    suggestion=$(echo "$suggestion" | grep -v '^[[:space:]]*$' | head -n 1)
    # Remove leading and trailing backticks if present
    if [[ "$suggestion" == \`* ]]; then
       suggestion="${suggestion#\`}"
    fi
    if [[ "$suggestion" == *\` ]]; then
       suggestion="${suggestion%\`}"
    fi
    BUFFER="$suggestion"
    zle end-of-line
  fi
}

zle -N _llm_autosuggest
bindkey '^Z' _llm_autosuggest
# === LLM MLX Autosuggestion Widget End ===

EOF

echo "Installation complete. Please restart your terminal to apply the changes."