# koh3.plugin.zsh — Koh 3 in your zsh prompt, plus tab completion.
#
#   kiln · v2.1.0 · 6 saves · 4.2 MB !
#   →
#
# The segment reads `koh status --json`, never the .koh file itself, so it
# works with any schema and needs no sqlite3. `!` means unsaved changes and
# `⚑` a flagged head save. Koh 2 repos (a .koh folder) are left to Koh 2's
# plugin.
#
# Install: source this file from ~/.zshrc. To keep your own prompt, set
# KOH_DEFAULT_PROMPT=false first and put $(koh_prompt) where you want it.
# koh_prompt_info gives the same text without colour, for Starship or p10k:
#
#   [custom.koh]
#   command = "koh_prompt_info"
#   when = "test -f .koh"
#   shell = ["zsh"]

: ${KOH_COLOR_DIR:=73}
: ${KOH_COLOR_VERSION:=214}
: ${KOH_COLOR_META:=245}
: ${KOH_COLOR_ALERT:=203}
: ${KOH_SHOW_VERSION:=true}
: ${KOH_SHOW_SAVES:=true}
: ${KOH_SHOW_SIZE:=true}
: ${KOH_PROMPT_CHAR:=→}
: ${KOH_DEFAULT_PROMPT:=true}

_koh_bin() {
  if [[ -n $KOH_BIN ]]; then print -r -- $KOH_BIN
  else print koh
  fi
}

# The folder holding a Koh 3 .koh file, searching upward like koh does.
_koh_root() {
  local dir=$PWD
  while true; do
    [[ -f $dir/.koh ]] && { print -r -- $dir; return 0; }
    [[ -d $dir/.koh || $dir == / ]] && return 1
    dir=${dir:h}
  done
}

# The project's own version, from whichever manifest it keeps.
_koh_version() {
  local root=$1 line
  if [[ -f $root/VERSION ]]; then
    read -r line <$root/VERSION && print -r -- ${line// /}
  elif [[ -f $root/v.mod ]]; then
    [[ "$(<$root/v.mod)" =~ "version: *'([^']+)'" ]] && print -r -- $match[1]
  elif [[ -f $root/package.json ]]; then
    [[ "$(<$root/package.json)" =~ '"version" *: *"([^"]+)"' ]] && print -r -- $match[1]
  elif [[ -f $root/Cargo.toml ]]; then
    [[ "$(<$root/Cargo.toml)" =~ $'(^|\n)version *= *"([^"]+)"' ]] && print -r -- $match[2]
  elif [[ -f $root/build.zig.zon ]]; then
    [[ "$(<$root/build.zig.zon)" =~ '\.version *= *"([^"]+)"' ]] && print -r -- $match[1]
  elif [[ -f $root/pyproject.toml ]]; then
    [[ "$(<$root/pyproject.toml)" =~ $'(^|\n)version *= *"([^"]+)"' ]] && print -r -- $match[2]
  fi
}

_koh_size() {
  local -F value=$1
  local unit=B
  for next in KB MB GB; do
    (( value < 1024 )) && break
    (( value /= 1024 ))
    unit=$next
  done
  [[ $unit == B ]] && print -r -- "$1 B" || printf '%.1f %s' $value $unit
}

# _koh_segment <colour?> — the segment text, or nothing outside a Koh 3 repo.
_koh_segment() {
  local colour=$1 root json
  root=$(_koh_root) || return
  json=$(cd $root && command $(_koh_bin) status --json 2>/dev/null) || return
  local saves size dirty flagged version
  [[ $json =~ '"saves":([0-9]+)' ]] && saves=$match[1]
  [[ $json =~ '"size_bytes":([0-9]+)' ]] && size=$(_koh_size $match[1])
  [[ $json == *'"dirty":true'* ]] && dirty=1
  [[ $json == *'"head_flagged":true'* ]] && flagged=1
  [[ $KOH_SHOW_VERSION == true ]] && version=$(_koh_version $root)

  local c_dir= c_version= c_meta= c_alert= reset=
  if [[ -n $colour ]]; then
    c_dir="%F{$KOH_COLOR_DIR}" c_version="%F{$KOH_COLOR_VERSION}"
    c_meta="%F{$KOH_COLOR_META}" c_alert="%F{$KOH_COLOR_ALERT}" reset=%f
  fi
  local out="${c_dir}${root:t}${reset}"
  [[ -n $version ]] && out+="${c_meta} · ${reset}${c_version}v${version#v}${reset}"
  [[ $KOH_SHOW_SAVES == true && -n $saves ]] && out+="${c_meta} · $saves save$( (( saves == 1 )) || print s)${reset}"
  [[ $KOH_SHOW_SIZE == true && -n $size ]] && out+="${c_meta} · $size${reset}"
  [[ -n $dirty ]] && out+=" ${c_alert}!${reset}"
  [[ -n $flagged ]] && out+=" ${c_alert}⚑${reset}"
  print -rn -- $out
}

koh_prompt() { _koh_segment colour; }
koh_prompt_info() { _koh_segment; }

if [[ $KOH_DEFAULT_PROMPT == true ]]; then
  setopt prompt_subst
  _koh_precmd() {
    local last=$?
    local segment=$(koh_prompt)
    if [[ -z $segment ]]; then
      segment="%F{$KOH_COLOR_DIR}%2~%f"
    fi
    local char="%F{$KOH_COLOR_META}${KOH_PROMPT_CHAR}%f"
    (( last != 0 )) && char="%F{$KOH_COLOR_ALERT}${KOH_PROMPT_CHAR}%f"
    PROMPT="${segment}"$'\n'"${char} "
  }
  autoload -Uz add-zsh-hook
  add-zsh-hook precmd _koh_precmd
fi

# Tab completion: the commands of `koh`'s usage, and save IDs where one is taken.
_koh3() {
  local -a commands=(
    'save:snapshot the working tree'
    'checkpoint:save, even with no changes'
    'status:what changed since the last save'
    'view:the history, saves and checkpoints'
    'diff:unified diff'
    'undo:step back n saves'
    'revert:discard unsaved changes'
    'load:restore any save'
    'remote:choose S3 or a folder, once per machine'
    'offer:send new saves to the remote'
    'steal:bring in saves from the remote'
    'flag:mark a save unsafe'
    'pluck:apply one save, or its inverse'
    'peek:restore a save into a temp folder'
    'init:create .koh here'
    'doctor:check repository health'
    'ignore:ignore or unignore files and folders'
    'strip:remove a path from history'
    'build:build from koh.build'
    'update:update Koh itself'
    'migrate:convert a Koh 2 repo here'
    'man:the manual, built in'
  )
  _koh3_ids() {
    local -a ids
    ids=(${(f)"$(command $(_koh_bin) view --json 2>/dev/null |
      grep -o '"id":"[0-9a-f]*","parent":"[0-9a-f]*","message":"[^"]*"' |
      sed -E 's/"id":"([0-9a-f]*)".*"message":"([^"]*)"/\1:\2/' | head -30)"})
    _describe 'save' ids
  }
  if (( CURRENT == 2 )); then
    _describe 'koh command' commands
    return
  fi
  case $words[2] in
    save|checkpoint) _arguments '--agent[model name]:agent:' '--session[session name]:session:' '--json' ;;
    view)            _arguments '--checkpoints[start on the checkpoints tab]' '--since[filter by time]:expr:' '--session[session name]:session:' '--json' ;;
    diff|load|peek)  _koh3_ids ;;
    pluck)           _arguments '--reverse' '--dry-run' '*:save:_koh3_ids' ;;
    flag)            _arguments '--reason[why]:reason:' '--clear' '--json' '*:save:_koh3_ids' ;;
    undo)            _arguments '--session' '--yes' ;;
    revert)          _arguments '--yes' ;;
    remote)          _arguments '1:kind:(folder s3)' '--forget' '--yes' '--json' ;;
    offer)           _arguments '--yes' '--builds' '--notes:notes:' '--json' ;;
    steal)           _arguments '--yes' '--list' '--mine' '--theirs' '--json' ;;
    ignore)          _arguments '--global[edit ~/.kohignore for every repo]' '--yes[skip the question]' '--json' '*:path:_files' ;;
    strip)           _arguments '--yes[skip the yes/no question; paths still typed]' '--save[one save]:save:_koh3_ids' '*:path:_files' ;;
    doctor)          _arguments '--fast[skip the slow remote audit]' '--json[print the report as JSON]' '--yes[remove old peek folders]' '--backup[copy .koh to]:path:_files' ;;
    build)           _arguments '*--platform[target platform]:platform:(macos-aarch64 macos-x86_64 linux-aarch64 linux-x86_64)' '--all' ;;
    update)          _arguments '--check' '--yes' '--json' ;;
    migrate)         _arguments '--lane[Koh 2 lane]:lane:(dev main)' '*--with[a lost version, from a backup]:file:_files' '--yes[skip the question]' '--json' ;;
    man)             _describe 'command' commands ;;
  esac
}
if (( $+functions[compdef] )); then
  compdef _koh3 ${$(_koh_bin):t}
fi