#!/bin/bash
# Rewind workspace files to a specific turn snapshot.
# Resets each repository to committed state, then reapplies the turn patch.

set -o pipefail

TASK_OUTPUT_DIR="${TASK_OUTPUT_DIR:-/task-output}"
WORKSPACE_DIR="${WORKSPACE_DIR:-/workspace}"
TARGET_TURN="${TARGET_TURN:-}"
turn_dir="$TASK_OUTPUT_DIR/patches/turn-$TARGET_TURN"

if ! [[ "$TARGET_TURN" =~ ^[0-9]+$ ]]; then
  echo "{\"target_turn\":0,\"success\":false,\"error\":\"TARGET_TURN must be a non-negative integer\",\"repos\":[]}"
  exit 1
fi

temp_dir=$(mktemp -d)
trap "rm -rf $temp_dir" EXIT

overall_success=true
repo_index=0

emit_repo_result() {
  local output_file="$1"
  local repo_name="$2"
  local repo_path="$3"
  local has_patch="$4"
  local applied="$5"
  local success="$6"
  local error_message="$7"
  local files_changed="$8"
  local lines_added="$9"
  local lines_deleted="${10}"

  jq -n \
    --arg name "$repo_name" \
    --arg path "$repo_path" \
    --argjson has_patch "$has_patch" \
    --argjson applied "$applied" \
    --argjson success "$success" \
    --arg error "$error_message" \
    --argjson files_changed "$files_changed" \
    --argjson lines_added "$lines_added" \
    --argjson lines_deleted "$lines_deleted" \
    '{
      name: $name,
      path: $path,
      has_patch: $has_patch,
      applied: $applied,
      success: $success,
      error: ($error | if . == "" then null else . end),
      files_changed: $files_changed,
      lines_added: $lines_added,
      lines_deleted: $lines_deleted
    }' > "$output_file"
}

process_repo() {
  local repo_name="$1"
  local repo_path="$2"
  local output_file="$3"
  local patch_file="$turn_dir/$repo_name.patch"

  local has_patch=false
  local applied=false
  local success=true
  local error_message=""
  local files_changed=0
  local lines_added=0
  local lines_deleted=0

  if [ ! -d "$repo_path/.git" ]; then
    success=false
    error_message="Repository is missing .git directory"
    emit_repo_result "$output_file" "$repo_name" "$repo_path" "$has_patch" "$applied" "$success" "$error_message" "$files_changed" "$lines_added" "$lines_deleted"
    return 1
  fi

  cd "$repo_path" || {
    success=false
    error_message="Failed to access repository directory"
    emit_repo_result "$output_file" "$repo_name" "$repo_path" "$has_patch" "$applied" "$success" "$error_message" "$files_changed" "$lines_added" "$lines_deleted"
    return 1
  }

  if ! git reset --hard HEAD >/dev/null 2>&1; then
    success=false
    error_message="git reset --hard failed"
  fi

  if $success && ! git clean -fd >/dev/null 2>&1; then
    success=false
    error_message="git clean -fd failed"
  fi

  if $success && [ -f "$patch_file" ] && [ -s "$patch_file" ]; then
    has_patch=true
    if git apply --whitespace=nowarn "$patch_file" >/dev/null 2>&1; then
      applied=true
    else
      success=false
      error_message="git apply failed for turn patch"
    fi
  fi

  if $success && ! git add -A >/dev/null 2>&1; then
    success=false
    error_message="git add -A failed"
  fi

  if $success; then
    local stat_output
    stat_output=$(git diff --cached --stat 2>/dev/null | tail -1)
    if [ -n "$stat_output" ]; then
      files_changed=$(echo "$stat_output" | grep -oE '[0-9]+ file' | grep -oE '[0-9]+' || echo "0")
      lines_added=$(echo "$stat_output" | grep -oE '[0-9]+ insertion' | grep -oE '[0-9]+' || echo "0")
      lines_deleted=$(echo "$stat_output" | grep -oE '[0-9]+ deletion' | grep -oE '[0-9]+' || echo "0")
    fi
    files_changed=${files_changed:-0}
    lines_added=${lines_added:-0}
    lines_deleted=${lines_deleted:-0}
  fi

  emit_repo_result "$output_file" "$repo_name" "$repo_path" "$has_patch" "$applied" "$success" "$error_message" "$files_changed" "$lines_added" "$lines_deleted"
  $success
}

process_repo_entry() {
  local repo_name="$1"
  local repo_path="$2"
  local output_file="$temp_dir/repo_$repo_index.json"

  process_repo "$repo_name" "$repo_path" "$output_file" || overall_success=false
  repo_index=$((repo_index + 1))
}

# Iterate configured repositories, or scan workspace if no config was provided.
if [ -n "$REPOS_CONFIG" ]; then
  repo_count=$(echo "$REPOS_CONFIG" | jq 'length' 2>/dev/null)
  if [ -n "$repo_count" ] && [ "$repo_count" -gt 0 ]; then
    for index in $(seq 0 $((repo_count - 1))); do
      repo_json=$(echo "$REPOS_CONFIG" | jq -c ".[$index]" 2>/dev/null)
      [ -z "$repo_json" ] || [ "$repo_json" = "null" ] && continue

      repo_name=$(echo "$repo_json" | jq -r 'if type == "object" then (.name // empty) else . end' 2>/dev/null)
      repo_path_rel=$(echo "$repo_json" | jq -r 'if type == "object" then (.path // empty) else empty end' 2>/dev/null)
      [ -z "$repo_name" ] || [ "$repo_name" = "null" ] && continue

      if [ -n "$repo_path_rel" ] && [ "$repo_path_rel" != "null" ]; then
        repo_path="$WORKSPACE_DIR/$repo_path_rel"
      else
        repo_path="$WORKSPACE_DIR/$repo_name"
      fi

      process_repo_entry "$repo_name" "$repo_path"
    done
  fi
else
  for repo_path in "$WORKSPACE_DIR"/*; do
    if [ -d "$repo_path/.git" ]; then
      repo_name=$(basename "$repo_path")
      process_repo_entry "$repo_name" "$repo_path"
    fi
  done
fi

# Collect repo results into one array
repos_json="["
first=true
for f in "$temp_dir"/repo_*.json; do
  [ -f "$f" ] && [ -s "$f" ] || continue
  if $first; then
    first=false
  else
    repos_json="$repos_json,"
  fi
  repos_json="$repos_json$(cat "$f")"
done
repos_json="$repos_json]"

if $overall_success; then
  overall_success_json=true
else
  overall_success_json=false
fi

jq -n \
  --argjson target_turn "$TARGET_TURN" \
  --argjson success "$overall_success_json" \
  --argjson repos "$repos_json" \
  '{
    target_turn: $target_turn,
    success: $success,
    repos: $repos
  }' > "$TASK_OUTPUT_DIR/rewind-result.json"

cat "$TASK_OUTPUT_DIR/rewind-result.json"
