nixos/home/common/optional/desktop/dwm/scripts/git-commit-ai.nix

{ pkgs, ... }:
{
  home.packages = [
    (pkgs.writers.writePython3Bin
      "tiktoken-cli"
      { libraries = [ pkgs.python3Packages.tiktoken ]; }
      ''
        import os
        import sys
        import tiktoken
      
        content = sys.stdin.read()
      
        default = "gpt-3.5-turbo-0301"
        model = os.getenv('TIKTOKEN_MODEL', default)
      
        try:
            encoding = tiktoken.encoding_for_model(model)
        except KeyError:
            encoding = tiktoken.get_encoding(default)
      
        print(len(encoding.encode(content)))
      '')

    (pkgs.writeShellScriptBin "git-commit-ai" ''
      # This bash script performs the following tasks:
      # 1. Takes a git diff of modifications made to a repo.
      # 2. Passes this diff to an LLM (Large Language Model) that generates a commit
      # message.
      #
      # Usage: ./git-commit-ai --git-diff-commands --model
      #
      #
      diff=$(git diff "$@" )
      [ -z "$diff" ] && echo "No diff provided. Did you stage the changes?" && exit 1
      tokens=$(echo "$diff" | tiktoken-cli)
      echo -e "\n\nEstimated tokens: $tokens\n\n"

      choose_model() {
          models=$(${pkgs.aichat}/bin/aichat --list-models)
          model=$(echo "$models" | fzf)
      }
      choose_model
      
      git config commit.template ".git/gitmessage"
      
      redo=""
      while true
      do
        
        [ -z "$redo" ] && echo "Using model: $model" && commit_message=$(${pkgs.aichat}/bin/aichat --model "$model" --role "git_commit" "$diff")
      
        echo ""
        echo "Message:"
        echo "---------------------------------------------"
        echo ""
        echo "$commit_message"
        echo ""
        echo "---------------------------------------------"
        echo "What would you like to do?"
        echo ""
        echo "Copy to clipboard (c)"
        echo "Git commit in editor (g)"
        echo "Open in editor (e)"
        echo "Show diff (d)"
        echo "Choose different model (m)"
        echo "Exit and do nothing (q)"
        
        read -r -p "Action: (c|g|e|q) " input
        
        case "$input" in
          "c") echo "$commit_message" | ${pkgs.xclip}/bin/xclip -i -selection CLIPBOARD && 
               ${pkgs.libnotify}/bin/notify-send "AI Buddy" "Copied to clipboard"
               exit 0
          ;;
          "g") echo "$commit_message" > ".git/gitmessage" &&
               git commit -e
               cat /dev/null > ".git/gitmessage" # clear the template
               exit 0
          ;;
        "d") echo -e "Current commit message:\n\n$commit_message\n\nDiff file:\n\n$diff" | less
               redo=true
               shift
          ;;
          "e") echo "$commit_message" | nvim -
               exit 0
          ;;
          "m") choose_model && shift
          ;;
          "q") exit 0
          ;;
            *) exit 0
        esac
      done  
      
    '')
  ];
}