#!/bin/bash
set -euo pipefail

REPO="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)"
ZSHRC="${ZDOTDIR:-$HOME}/.zshrc"
BEGIN_MARK="# >>> omarchy-zsh dev >>>"
END_MARK="# <<< omarchy-zsh dev <<<"
PKG_ZOPTIONS_LINE='[[ -f /usr/share/omarchy-zsh/shell/zoptions ]] && source /usr/share/omarchy-zsh/shell/zoptions'

usage() {
  cat <<USAGE
Usage: bin/dev-link [--omadots PATH]

Inject a marker block into ~/.zshrc that points fpath at this repo's
working tree and re-sources zoptions, so edits take effect in new shells
without rebuilding the package.

Options:
  --omadots PATH   Also overlay omadots' working tree (shell/aliases,
                   envs, functions, inits) so edits to the shared shell
                   config show up in dev shells. Defaults to
                   \$OMARCHY_DEV_OMADOTS_PATH if set.
USAGE
}

omadots=""
while (( $# )); do
  case "$1" in
    --omadots)
      omadots="${2:-}"
      shift 2 || { usage >&2; exit 2; }
      ;;
    --omadots=*)
      omadots="${1#--omadots=}"
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown argument: $1" >&2
      usage >&2
      exit 2
      ;;
  esac
done

# Flag wins; otherwise fall back to the env var.
omadots="${omadots:-${OMARCHY_DEV_OMADOTS_PATH:-}}"

if [[ -n $omadots ]]; then
  if [[ ! -d $omadots/config/shell ]]; then
    echo "omadots path does not look right: $omadots/config/shell missing" >&2
    exit 1
  fi
  omadots="$(cd -- "$omadots" && pwd)"
fi

if [[ ! -f $ZSHRC ]]; then
  echo "No ~/.zshrc found at $ZSHRC. Run omarchy-setup-zsh first." >&2
  exit 1
fi

# zsh has no auto-prepended user vendor dir like fish's vendor_*.d/.
# Instead of symlinking into /usr/share/zsh/site-functions/ (sudo), replace
# the packaged zoptions source line in ~/.zshrc with an idempotent dev block.
# This matters because zoptions can load plugins; sourcing the packaged
# zoptions first would leave removed plugins active before the dev copy loads.
{
  printf '%s\n' "$BEGIN_MARK"
  printf 'fpath=("%s/shell/completions" $fpath)\n' "$REPO"
  printf '[[ -f "%s/shell/zoptions" ]] && source "%s/shell/zoptions"\n' "$REPO" "$REPO"
  printf 'unfunction _omarchy 2>/dev/null\n'
  printf 'autoload -Uz _omarchy\n'
  if [[ -n $omadots ]]; then
    # Re-source the pure-data overlays in the same order as
    # /usr/share/omarchy-zsh/shell/all. We skip `inits` on purpose:
    # tool initializers (mise/starship/zoxide/try) aren't idempotent
    # under re-source, and the system version already ran. Edits to
    # inits need a real `bin/dev-pkg-test` rebuild to validate.
    for f in envs aliases functions; do
      printf '[[ -f "%s/config/shell/%s" ]] && source "%s/config/shell/%s"\n' \
        "$omadots" "$f" "$omadots" "$f"
    done
    printf 'print -P "%%F{yellow}⚠ omarchy-zsh dev linked: %%B%s%%b (omadots: %%B%s%%b) — run %%U'\''bin/dev-unlink'\''%%u to restore packaged version%%f"\n' \
      "$REPO" "$omadots"
  else
    printf 'print -P "%%F{yellow}⚠ omarchy-zsh dev linked: %%B%s%%b — run %%U'\''bin/dev-unlink'\''%%u to restore packaged version%%f"\n' \
      "$REPO"
  fi
  printf '%s\n' "$END_MARK"
} > "$ZSHRC.snippet"

# Always drop any existing block, then replace the packaged zoptions source
# line with the fresh dev block. If a custom ~/.zshrc no longer has the
# packaged source line, append the block as a fallback.
awk '
  /^# >>> omarchy-zsh dev/ { skip=1; next }
  skip && /^# <<< omarchy-zsh dev/ { skip=0; next }
  !skip { print }
' "$ZSHRC" > "$ZSHRC.tmp"

awk -v snippet="$ZSHRC.snippet" -v pkg_line="$PKG_ZOPTIONS_LINE" '
  $0 == pkg_line && !inserted {
    while ((getline line < snippet) > 0) print line
    close(snippet)
    inserted=1
    next
  }
  { print }
  END {
    if (!inserted) {
      print ""
      while ((getline line < snippet) > 0) print line
      close(snippet)
    }
  }
' "$ZSHRC.tmp" > "$ZSHRC"

rm -f "$ZSHRC.tmp" "$ZSHRC.snippet"

echo "Linked $REPO into $ZSHRC"
[[ -n $omadots ]] && echo "Overlaying omadots from $omadots"
echo "Open a new zsh shell to pick up the working-tree version."
