#!/usr/bin/env bash
# swarm-state: minimal state, task and lock manager for tmux AI swarm workflows.
set -euo pipefail

SWARM_DIR="${AUTO_TMUX_SWARM_DIR:-/tmp/ai_swarm}"
STATUS_LOG="$SWARM_DIR/status.log"
TASKS_TSV="$SWARM_DIR/tasks.tsv"
LOCKS_DIR="$SWARM_DIR/locks"
RESULTS_DIR="$SWARM_DIR/results"

usage() {
  cat <<'EOF'
swarm-state: tmux AI swarm state/task/lock helper

Usage:
  swarm-state.sh init [--dir DIR]
  swarm-state.sh log --target TARGET --status STATUS --message TEXT [--dir DIR]
  swarm-state.sh status [--dir DIR]
  swarm-state.sh task-add --text TEXT [--id ID] [--dir DIR]
  swarm-state.sh task-list [--dir DIR]
  swarm-state.sh task-claim --id ID --owner TARGET [--dir DIR]
  swarm-state.sh task-done --id ID --owner TARGET [--result TEXT] [--dir DIR]
  swarm-state.sh lock-acquire --name NAME --owner TARGET [--dir DIR]
  swarm-state.sh lock-release --name NAME --owner TARGET [--dir DIR] [--force]
  swarm-state.sh lock-list [--dir DIR]
  swarm-state.sh report [--dir DIR]

Environment:
  AUTO_TMUX_SWARM_DIR overrides the default /tmp/ai_swarm.
EOF
}

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

sanitize() {
  tr '\t\r\n' '   ' | sed -E 's/[[:space:]]+/ /g; s/^ //; s/ $//'
}

set_dir() {
  local dir="$1"
  SWARM_DIR="$dir"
  STATUS_LOG="$SWARM_DIR/status.log"
  TASKS_TSV="$SWARM_DIR/tasks.tsv"
  LOCKS_DIR="$SWARM_DIR/locks"
  RESULTS_DIR="$SWARM_DIR/results"
}

ensure_init() {
  mkdir -p "$LOCKS_DIR" "$RESULTS_DIR"
  touch "$STATUS_LOG" "$TASKS_TSV"
  if [[ ! -s "$TASKS_TSV" ]]; then
    printf 'id\tstatus\towner\ttext\tresult\n' > "$TASKS_TSV"
  fi
}

cmd_init() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      *) die "unknown init option: $1" ;;
    esac
  done
  ensure_init
  printf 'swarm initialized: %s\n' "$SWARM_DIR"
}

cmd_log() {
  local target="" status="" message=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      --target) target="${2:-}"; shift 2 ;;
      --status) status="${2:-}"; shift 2 ;;
      --message) message="${2:-}"; shift 2 ;;
      *) die "unknown log option: $1" ;;
    esac
  done
  [[ -n "$target" ]] || die "missing --target"
  [[ -n "$status" ]] || die "missing --status"
  [[ -n "$message" ]] || die "missing --message"
  ensure_init
  message="$(printf '%s' "$message" | sanitize)"
  printf '[%s] [%s] [%s] %s\n' "$(date +%H:%M:%S)" "$target" "$status" "$message" >> "$STATUS_LOG"
}

cmd_status() {
  local lines=30
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      -n|--lines) lines="${2:-}"; shift 2 ;;
      *) die "unknown status option: $1" ;;
    esac
  done
  ensure_init
  tail -n "$lines" "$STATUS_LOG"
}

cmd_task_add() {
  local id="" text=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      --id) id="${2:-}"; shift 2 ;;
      --text) text="${2:-}"; shift 2 ;;
      *) die "unknown task-add option: $1" ;;
    esac
  done
  [[ -n "$text" ]] || die "missing --text"
  ensure_init
  if [[ -z "$id" ]]; then
    id="task-$(date +%Y%m%d%H%M%S)-$$"
  fi
  text="$(printf '%s' "$text" | sanitize)"
  printf '%s\tTODO\t-\t%s\t-\n' "$id" "$text" >> "$TASKS_TSV"
  printf 'task added: %s\n' "$id"
}

cmd_task_list() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      *) die "unknown task-list option: $1" ;;
    esac
  done
  ensure_init
  column -t -s $'\t' "$TASKS_TSV" 2>/dev/null || cat "$TASKS_TSV"
}

update_task() {
  local id="$1" status="$2" owner="$3" result="$4"
  local tmp
  tmp="$(mktemp)"
  awk -F '\t' -v OFS='\t' -v id="$id" -v status="$status" -v owner="$owner" -v result="$result" '
    NR == 1 {print; next}
    $1 == id {$2 = status; $3 = owner; if (result != "") $5 = result; found = 1}
    {print}
    END {if (!found) exit 42}
  ' "$TASKS_TSV" > "$tmp" || {
    local code=$?
    rm -f "$tmp"
    if [[ "$code" == "42" ]]; then
      die "task not found: $id"
    fi
    exit "$code"
  }
  mv "$tmp" "$TASKS_TSV"
}

cmd_task_claim() {
  local id="" owner=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      --id) id="${2:-}"; shift 2 ;;
      --owner) owner="${2:-}"; shift 2 ;;
      *) die "unknown task-claim option: $1" ;;
    esac
  done
  [[ -n "$id" ]] || die "missing --id"
  [[ -n "$owner" ]] || die "missing --owner"
  ensure_init
  update_task "$id" "DOING" "$owner" ""
  cmd_log --target "$owner" --status "START" --message "claimed $id"
}

cmd_task_done() {
  local id="" owner="" result="done"
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      --id) id="${2:-}"; shift 2 ;;
      --owner) owner="${2:-}"; shift 2 ;;
      --result) result="${2:-}"; shift 2 ;;
      *) die "unknown task-done option: $1" ;;
    esac
  done
  [[ -n "$id" ]] || die "missing --id"
  [[ -n "$owner" ]] || die "missing --owner"
  ensure_init
  result="$(printf '%s' "$result" | sanitize)"
  update_task "$id" "DONE" "$owner" "$result"
  printf '%s\n' "$result" > "$RESULTS_DIR/$id.txt"
  cmd_log --target "$owner" --status "DONE" --message "$id $result"
}

lock_path() {
  local name="$1"
  name="$(printf '%s' "$name" | sed -E 's#[^A-Za-z0-9_.-]+#_#g')"
  printf '%s/%s.lock.d' "$LOCKS_DIR" "$name"
}

cmd_lock_acquire() {
  local name="" owner=""
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      --name) name="${2:-}"; shift 2 ;;
      --owner) owner="${2:-}"; shift 2 ;;
      *) die "unknown lock-acquire option: $1" ;;
    esac
  done
  [[ -n "$name" ]] || die "missing --name"
  [[ -n "$owner" ]] || die "missing --owner"
  ensure_init
  local path
  path="$(lock_path "$name")"
  if mkdir "$path" 2>/dev/null; then
    printf '%s\n' "$owner" > "$path/owner"
    printf '%s\n' "$(date -Is)" > "$path/created_at"
    cmd_log --target "$owner" --status "LOCK" --message "acquired $name"
    printf 'lock acquired: %s\n' "$name"
  else
    printf 'lock busy: %s owner=' "$name" >&2
    cat "$path/owner" >&2 2>/dev/null || printf 'unknown\n' >&2
    exit 2
  fi
}

cmd_lock_release() {
  local name="" owner="" force="0"
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      --name) name="${2:-}"; shift 2 ;;
      --owner) owner="${2:-}"; shift 2 ;;
      --force) force="1"; shift ;;
      *) die "unknown lock-release option: $1" ;;
    esac
  done
  [[ -n "$name" ]] || die "missing --name"
  [[ -n "$owner" ]] || die "missing --owner"
  ensure_init
  local path current_owner
  path="$(lock_path "$name")"
  [[ -d "$path" ]] || die "lock not found: $name"
  current_owner="$(cat "$path/owner" 2>/dev/null || true)"
  if [[ "$force" != "1" && "$current_owner" != "$owner" ]]; then
    die "lock owner mismatch: current=$current_owner requested=$owner"
  fi
  rm -rf "$path"
  cmd_log --target "$owner" --status "UNLOCK" --message "released $name"
  printf 'lock released: %s\n' "$name"
}

cmd_lock_list() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      *) die "unknown lock-list option: $1" ;;
    esac
  done
  ensure_init
  find "$LOCKS_DIR" -mindepth 1 -maxdepth 1 -type d -name '*.lock.d' -print | sort | while read -r path; do
    printf '%s\towner=%s\tcreated_at=%s\n' \
      "$(basename "$path" .lock.d)" \
      "$(cat "$path/owner" 2>/dev/null || printf unknown)" \
      "$(cat "$path/created_at" 2>/dev/null || printf unknown)"
  done
}

cmd_report() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --dir) set_dir "${2:-}"; shift 2 ;;
      *) die "unknown report option: $1" ;;
    esac
  done
  ensure_init
  printf '# swarm report\n\n'
  printf 'dir: %s\n\n' "$SWARM_DIR"
  printf '## tasks\n\n'
  cmd_task_list
  printf '\n## locks\n\n'
  cmd_lock_list || true
  printf '\n## recent status\n\n'
  cmd_status -n 20
}

main() {
  local cmd="${1:-help}"
  shift || true
  case "$cmd" in
    help|-h|--help) usage ;;
    init) cmd_init "$@" ;;
    log) cmd_log "$@" ;;
    status) cmd_status "$@" ;;
    task-add) cmd_task_add "$@" ;;
    task-list) cmd_task_list "$@" ;;
    task-claim) cmd_task_claim "$@" ;;
    task-done) cmd_task_done "$@" ;;
    lock-acquire) cmd_lock_acquire "$@" ;;
    lock-release) cmd_lock_release "$@" ;;
    lock-list) cmd_lock_list "$@" ;;
    report) cmd_report "$@" ;;
    *) die "unknown command: $cmd" ;;
  esac
}

main "$@"