#!/usr/bin/env sh
set -eu

CODEX_HOME_DIR="${CODEX_HOME:-${HOME}/.codex}"
AGENTS_MD="${CODEX_HOME_DIR}/AGENTS.md"
BACKUP_MD="${AGENTS_MD}.pre-zg-force.bak"
START_MARKER='<!-- ZG_FORCE_START v1 -->'
START_PREFIX='<!-- ZG_FORCE_START'
END_MARKER='<!-- ZG_FORCE_END -->'

usage() {
  cat <<'USAGE'
Usage: install-zg-codex.sh [install|uninstall|uninstall-all|status]

Commands:
  install        Install zvec-grep if needed, configure its Codex MCP integration,
                 and add mandatory zg workspace-search rules to global AGENTS.md.
  uninstall      Remove only this script's ZG_FORCE block from global AGENTS.md.
  uninstall-all  Remove the ZG_FORCE block and run the official zg Codex uninstall.
  status         Show zvec-grep and global AGENTS.md integration status.

Environment:
  CODEX_HOME     Codex config directory. Defaults to ~/.codex.
USAGE
}

die() {
  printf 'error: %s\n' "$*" >&2
  exit 1
}

info() {
  printf '%s\n' "$*"
}

count_line_prefix() {
  file=$1
  prefix=$2
  if [ ! -f "$file" ]; then
    printf '0\n'
    return
  fi
  awk -v prefix="$prefix" 'index($0, prefix) == 1 { count++ } END { print count + 0 }' "$file"
}

validate_managed_block() {
  starts=$(count_line_prefix "$AGENTS_MD" "$START_PREFIX")
  ends=$(count_line_prefix "$AGENTS_MD" "$END_MARKER")

  [ "$starts" -le 1 ] || die "multiple ZG_FORCE start markers found in $AGENTS_MD"
  [ "$ends" -le 1 ] || die "multiple ZG_FORCE end markers found in $AGENTS_MD"
  [ "$starts" -eq "$ends" ] || die "malformed ZG_FORCE block in $AGENTS_MD"
}

remove_force_block() {
  [ -f "$AGENTS_MD" ] || return 0
  validate_managed_block

  starts=$(count_line_prefix "$AGENTS_MD" "$START_PREFIX")
  [ "$starts" -eq 1 ] || return 0

  tmp=$(mktemp "${CODEX_HOME_DIR}/.agents-zg.XXXXXX")
  awk -v start="$START_PREFIX" -v end="$END_MARKER" '
    index($0, start) == 1 { skipping=1; next }
    skipping && index($0, end) == 1 { skipping=0; next }
    !skipping { print }
  ' "$AGENTS_MD" > "$tmp"

  # Trim only trailing blank lines introduced by removal while preserving other content.
  awk '
    { lines[NR]=$0 }
    END {
      last=NR
      while (last > 0 && lines[last] == "") last--
      for (i=1; i<=last; i++) print lines[i]
    }
  ' "$tmp" > "${tmp}.trim"
  cat "${tmp}.trim" > "$AGENTS_MD"
  rm -f "$tmp" "${tmp}.trim"
}

append_force_block() {
  mkdir -p "$CODEX_HOME_DIR"
  touch "$AGENTS_MD"
  validate_managed_block
  remove_force_block

  if [ -s "$AGENTS_MD" ]; then
    printf '\n\n' >> "$AGENTS_MD"
  fi

  cat >> "$AGENTS_MD" <<'RULES'
<!-- ZG_FORCE_START v1 -->
# ZG (zvec-grep) - Mandatory Workspace Search

## Rule

For repository or workspace content search, use zvec-grep (`zg`) as the default and mandatory search layer.

- For semantic, conceptual, fuzzy, relationship, chronology, causality, comparison, or cross-file discovery, use the `zvec_grep_search` MCP tool first. If the MCP tool is unavailable, use `zg query "<natural-language query>"`.
- For exact literals, identifiers, strings, regexes, or exhaustive occurrence checks, use `zg query --rg ...` instead of raw `rg`, `grep`, or `git grep`.
- For known terms that need ranked lexical retrieval, use `zg query --fts "<terms>"`.
- After semantic discovery reveals an exact identifier or string, verify exact occurrences with `zg query --rg ...` when verification is useful.
- Scope broad searches with paths, `-g/--glob`, `-t/--type`, or `-T/--type-not` whenever possible.

Do not use raw `rg`, `grep`, `git grep`, or equivalent native content-search tools for workspace search unless the user explicitly requests them or `zg` is unavailable or fails. If a fallback is necessary, state the reason briefly and continue with the narrowest suitable fallback.

This rule applies to searching file contents and discovering relevant workspace evidence. Normal filesystem operations for already-known paths, directory listing, existence checks, and reading a known file may use the ordinary filesystem tools.

## Index discipline

`zvec_grep_search` and indexed `zg query` routes require a workspace index. Do not create, rebuild, change the embedding model of, or delete a persistent index without explicit user authorization.

If indexed search is required and the index is unavailable:

1. Run `zg status` only when needed to diagnose the missing or stale index.
2. Use `zg query --rg ...` if an exact-search route can still answer the question without an index.
3. Otherwise report that indexed retrieval requires `zg index` and request authorization before creating or rebuilding it.

Do not run a status preflight before every search.

## Examples

```bash
# Semantic / hybrid discovery
zg query "where authentication state is restored"

# Ranked lexical retrieval
zg query --fts "AuthService ForbiddenError"

# Exact literal / identifier search
zg query --rg -n -F "AuthService" src

# Regex with scope
zg query --rg -n 'TODO|FIXME' -g '*.kt' src
```
<!-- ZG_FORCE_END -->
RULES
}

ensure_node_22() {
  command -v node >/dev/null 2>&1 || die "Node.js 22+ is required to install zvec-grep"
  major=$(node -p 'process.versions.node.split(".")[0]' 2>/dev/null || true)
  case "$major" in
    ''|*[!0-9]*) die "could not determine Node.js version" ;;
  esac
  [ "$major" -ge 22 ] || die "Node.js 22+ is required; found $(node --version 2>/dev/null || printf unknown)"
}

ensure_zg() {
  if command -v zg >/dev/null 2>&1; then
    return 0
  fi

  command -v npm >/dev/null 2>&1 || die "npm is required to install @zvec/zvec-grep"
  ensure_node_22
  info "Installing @zvec/zvec-grep globally..."
  npm install -g @zvec/zvec-grep
  command -v zg >/dev/null 2>&1 || die "zg was installed but is not available on PATH"
}

backup_agents_once() {
  mkdir -p "$CODEX_HOME_DIR"
  if [ -f "$AGENTS_MD" ] && [ ! -e "$BACKUP_MD" ]; then
    cp "$AGENTS_MD" "$BACKUP_MD"
    info "Backup: $BACKUP_MD"
  fi
}

install_all() {
  ensure_zg
  backup_agents_once

  info "Configuring official zvec-grep Codex integration..."
  zg install --target codex --mcp-toolset agent --yes

  append_force_block

  info "Installed mandatory zg search policy: $AGENTS_MD"
  info "zvec-grep: $(zg version 2>/dev/null || zg --version 2>/dev/null || printf 'installed')"
  info "Open a new Codex session for the global instructions and MCP configuration to take effect."
}

uninstall_force() {
  remove_force_block
  info "Removed ZG_FORCE policy from: $AGENTS_MD"
}

uninstall_all() {
  uninstall_force
  if command -v zg >/dev/null 2>&1; then
    zg uninstall --target codex --yes
    info "Removed official zvec-grep Codex integration."
  else
    info "zg is not on PATH; skipped official zvec-grep uninstall."
  fi
}

show_status() {
  printf 'CODEX_HOME: %s\n' "$CODEX_HOME_DIR"
  printf 'AGENTS.md:  %s\n' "$AGENTS_MD"

  if command -v zg >/dev/null 2>&1; then
    printf 'zg:         %s\n' "$(command -v zg)"
    printf 'zg version: %s\n' "$(zg version 2>/dev/null || zg --version 2>/dev/null || printf unknown)"
  else
    printf 'zg:         not found\n'
  fi

  if [ "$(count_line_prefix "$AGENTS_MD" "$START_PREFIX")" -eq 1 ]; then
    printf 'ZG_FORCE:   installed\n'
  else
    printf 'ZG_FORCE:   not installed\n'
  fi

  if [ -f "$AGENTS_MD" ] && grep -q 'ZVEC_GREP_START' "$AGENTS_MD" 2>/dev/null; then
    printf 'official:   zvec-grep managed AGENTS.md block detected\n'
  else
    printf 'official:   managed AGENTS.md block not detected\n'
  fi
}

command=${1:-install}
case "$command" in
  install) install_all ;;
  uninstall) uninstall_force ;;
  uninstall-all) uninstall_all ;;
  status) show_status ;;
  -h|--help|help) usage ;;
  *) usage >&2; die "unknown command: $command" ;;
esac