#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
LB_DIR="$ROOT_DIR/.loop-build"
STATE_FILE="$LB_DIR/state/current_task.json"
HISTORY_FILE="$LB_DIR/state/history.ndjson"
POLICY_DIR="$LB_DIR/config"
POLICY_ENV="$POLICY_DIR/policy.env"
ALLOWLIST="$POLICY_DIR/allowlist.txt"
DENYLIST="$POLICY_DIR/denylist.txt"
SCRIPTS_DIR="$LB_DIR/scripts"

if [[ ! -f "$POLICY_ENV" ]]; then
  echo "[loop-build] missing policy file: $POLICY_ENV"
  exit 1
fi
source "$POLICY_ENV"

if ! command -v jq >/dev/null 2>&1; then
  echo "[loop-build] missing dependency: jq"
  exit 1
fi

if [[ ! -f "$STATE_FILE" ]]; then
  echo "[loop-build] missing state file: $STATE_FILE"
  exit 1
fi

now_ts() {
  date -u +%Y-%m-%dT%H:%M:%SZ
}

state_get() {
  jq -r "$1" "$STATE_FILE"
}

state_set_expr() {
  local expr="$1"
  local tmp
  tmp=$(mktemp)
  jq "$expr" "$STATE_FILE" > "$tmp"
  mv "$tmp" "$STATE_FILE"
}

state_set_string() {
  local path="$1"
  local value="$2"
  local tmp
  tmp=$(mktemp)
  jq --arg value "$value" "${path} = \$value" "$STATE_FILE" > "$tmp"
  mv "$tmp" "$STATE_FILE"
}

state_set_number() {
  local path="$1"
  local value="$2"
  local tmp
  tmp=$(mktemp)
  jq "${path} = ${value}" "$STATE_FILE" > "$tmp"
  mv "$tmp" "$STATE_FILE"
}

state_append() {
  local path="$1"
  local value="$2"
  local tmp
  tmp=$(mktemp)
  jq --arg value "$value" "${path} += [\$value]" "$STATE_FILE" > "$tmp"
  mv "$tmp" "$STATE_FILE"
}

state_set_last_verification() {
  local target="$1"
  local status="$2"
  local tmp
  tmp=$(mktemp)
  jq --arg target "$target" --arg status "$status" \
    '.last_verification = {"target": $target, "status": $status}' \
    "$STATE_FILE" > "$tmp"
  mv "$tmp" "$STATE_FILE"
}

state_mark_updated() {
  state_set_expr '.updated_at = "'"$(now_ts)"'"'
}

touch_history() {
  [[ -f "$HISTORY_FILE" ]] || : > "$HISTORY_FILE"
}

append_history() {
  local step_id="$1"
  local result="$2"
  local summary="$3"
  local verification="$4"
  local next_recommendation="$5"
  touch_history
  jq -n \
    --arg ts "$(now_ts)" \
    --arg sid "$step_id" \
    --arg res "$result" \
    --arg sum "$summary" \
    --arg v "$verification" \
    --arg next "$next_recommendation" \
    '{"timestamp":$ts,"step_id":$sid,"result":$res,"summary":$sum,"verification_result":$v,"next_step_recommendation":$next}' \
    >> "$HISTORY_FILE"
  echo >> "$HISTORY_FILE"
}

match_prefix() {
  local token="$1" list="$2"
  local pattern
  [[ -f "$list" ]] || return 1
  while IFS= read -r pattern; do
    [[ -z "$pattern" || "$pattern" =~ ^# ]] && continue
    if [[ "$token" == "$pattern" || "$token" == "$pattern"* ]]; then
      return 0
    fi
  done < "$list"
  return 1
}

policy_unlock_key() {
  local cmd="$1"
  local base="$2"

  if [[ "$base" == "sudo" ]]; then
    echo "ALLOW_SUDO"
    return
  fi
  if [[ "$base" == "rm" && "$cmd" == *" -rf"* ]]; then
    echo "ALLOW_RM_RF"
    return
  fi
  if [[ "$base" == "curl" || "$base" == "wget" ]]; then
    if [[ "$cmd" == *"|"* ]]; then
      echo "ALLOW_CURL_BASH"
      return
    fi
    echo "ALLOW_NETWORK"
    return
  fi
  if [[ "$base" == "git" && ("$cmd" == *" clone "* || "$cmd" == *" push "* || "$cmd" == *" pull "*) ]]; then
    echo "ALLOW_GIT_NETWORK"
    return
  fi
  if [[ "$base" == npm || "$base" == yarn || "$base" == pnpm || "$base" == pip || "$base" == pip3 || "$base" == composer ]]; then
    if [[ "$cmd" == *" install"* || "$cmd" == " install"* || "$cmd" == *" i "* || "$cmd" == " i "* ]]; then
      echo "ALLOW_INSTALL"
      return
    fi
  fi

  echo ""
}

policy_check() {
  local cmd="$1"
  local task_context="$2" # 1 for approved task step
  cmd="$(printf '%s' "$cmd" | sed 's/^ *//;s/ *$//')"
  local segments
  local segment
  local token
  local key

  segments="$(printf '%s' "$cmd" | sed -E 's/&&/\n/g; s/\|\|/\n/g; s/[;|]/\n/g')"
  while IFS= read -r segment; do
    segment="$(printf '%s' "$segment" | sed 's/^ *//;s/ *$//')"
    [[ -z "$segment" ]] && continue
    token="${segment%% *}"
    token="${token##*/}"

    if match_prefix "$segment" "$DENYLIST" || match_prefix "$token" "$DENYLIST"; then
      key="$(policy_unlock_key "$segment" "$token")"
      if [[ -n "$key" && "${!key:-0}" == "1" && "${POLICY_SECONDARY_CONFIRM:-0}" == "1" ]]; then
        continue
      fi
      echo "[policy] blocked by denylist segment: $segment"
      [[ -n "$key" ]] && echo "unlock: set $key=1 and POLICY_SECONDARY_CONFIRM=1 in $POLICY_ENV"
      return 2
    fi

    if match_prefix "$segment" "$ALLOWLIST" || match_prefix "$token" "$ALLOWLIST"; then
      continue
    fi

    if [[ "$task_context" != "1" ]]; then
      echo "[policy] blocked: this command requires approval context"
      echo "segment: $segment"
      return 2
    fi
  done <<< "$segments"

  return 0
}

error_excerpt() {
  local file="$1"
  local hit
  if [[ ! -s "$file" ]]; then
    return
  fi
  if grep -nEi "error|fail|fatal|exception|traceback" "$file" >/dev/null 2>&1; then
    hit=$(grep -nEi "error|fail|fatal|exception|traceback" "$file" | tail -n 1 | cut -d: -f1)
    awk -v start=$((hit > 50 ? hit - 50 : 1)) -v end=$((hit + 50)) 'NR>=start && NR<=end { print NR":"$0 }' "$file"
  else
    tail -n 80 "$file"
  fi
}

run_command() {
  local cmd="$1"
  local ctx="$2"
  local out err rc
  local out_file err_file
  out_file=$(mktemp)
  err_file=$(mktemp)

  if ! policy_check "$cmd" "$ctx"; then
    rm -f "$out_file" "$err_file"
    return 2
  fi

  set +e
  bash -lc "$cmd" >"$out_file" 2>"$err_file"
  rc=$?
  set -e

  if (( rc != 0 )); then
    echo "[command][FAIL] $cmd"
    error_excerpt "$err_file"
    rm -f "$out_file" "$err_file"
    return "$rc"
  fi

  if [[ -s "$out_file" ]]; then
    echo "[command][OUT]"
    tail -n 40 "$out_file"
  fi

  rm -f "$out_file" "$err_file"
  return 0
}

collect_changed_files() {
  if [[ ! -d "$ROOT_DIR/.git" ]]; then
    return 0
  fi
  git -C "$ROOT_DIR" status --short --untracked-files=normal | awk '{print $NF}' | sort -u
}

count_lines() {
  local file="$1"
  [[ -f "$file" ]] || { echo 0; return; }
  wc -l < "$file" | tr -d ' '
}

collect_step_actions() {
  local step_json="$1"
  local type
  type="$(jq -r 'type' <<<"$step_json")"
  case "$type" in
    string)
      jq -r '.' <<<"$step_json"
      ;;
    array)
      jq -r '.[]' <<<"$step_json"
      ;;
    object)
      if jq -e '.command' <<<"$step_json" >/dev/null 2>&1; then
        jq -r '.command // empty' <<<"$step_json"
      elif jq -e '.commands' <<<"$step_json" >/dev/null 2>&1; then
        jq -r '.commands[]?' <<<"$step_json"
      else
        return 1
      fi
      ;;
    *)
      return 1
      ;;
  esac
}

print_plan() {
  local steps status pending next
  status="$(state_get '.plan_status')"
  pending="$(state_get '.pending_action')"
  next="$(state_get '.next_step')"
  steps="$(state_get '.plan_steps | length')"

  echo "task_id: $(state_get '.task_id')"
  echo "plan_status: $status"
  echo "pending_action: $pending"
  echo "next_step: $next"
  echo "total_steps: $steps"
  echo "required_files: $(state_get '.required_files | join(",")')"
  echo "verify_targets: $(state_get '.verify_targets | join(",")')"
  echo "completion: $(state_get '.completed_steps')"

  if (( steps == 0 )); then
    echo "No plan present"
    return
  fi

  for i in $(seq 0 $((steps - 1))); do
    sid=$(jq -r ".plan_steps[$i].step_id // ( $i + 1 )" "$STATE_FILE")
    risk=$(jq -r ".plan_steps[$i].risk_level // \"LOW\"" "$STATE_FILE")
    verify=$(jq -r ".plan_steps[$i].verification // \"\"" "$STATE_FILE")
    action=$(jq -r ".plan_steps[$i].action // \"\"" "$STATE_FILE")
    rollback=$(jq -r ".plan_steps[$i].rollback // \"\"" "$STATE_FILE")
    echo "- step=$sid risk=$risk verify=$verify"
    echo "  action: $action"
    echo "  rollback: $rollback"
  done
}

status() {
  print_plan
  if [[ "$(state_get '.pending_action')" == "need_plan_approval" ]]; then
    echo "approval: APPROVE_PLAN / REVISE_PLAN / CANCEL_TASK"
  elif [[ "$(state_get '.pending_action')" == "step_ready" ]]; then
    echo "approval: step N then APPROVE_NEXT / REVISE_PLAN / CANCEL_TASK"
  elif [[ "$(state_get '.pending_action')" == "await_next_approval" ]]; then
    echo "approval: APPROVE_NEXT / REVISE_PLAN / CANCEL_TASK"
  fi
}

ensure_plan_exists() {
  local steps
  steps="$(state_get '.plan_steps | length')"
  if [[ "$steps" == "0" ]]; then
    echo "[plan] no plan_steps found; generate with planner first."
    exit 1
  fi
}

ensure_open_json() {
  # Ensure required state fields exist for operations in malformed initial states.
  state_set_expr '.completed_steps |= if type=="array" then . else [] end'
  state_set_expr '.open_questions |= if type=="array" then . else [] end'
}

run_verify_if_needed() {
  local target="$1"
  if [[ -z "$target" || "$target" == "null" ]]; then
    echo "not_run"
    return 0
  fi

  if ! policy_check "${SCRIPTS_DIR}/verify.sh $target" 1 >/dev/stderr; then
    echo "blocked"
    return 0
  fi

  local out err rc
  out=$(mktemp)
  err=$(mktemp)
  set +e
  APPROVE=1 "$SCRIPTS_DIR/verify.sh" "$target" >"$out" 2>"$err"
  rc=$?
  set -e

  if (( rc != 0 )); then
    echo "FAIL"
    error_excerpt "$err" >&2
    rm -f "$out" "$err"
    return 0
  fi

  if [[ -s "$out" ]]; then
    head -n 20 "$out" >&2
  fi
  rm -f "$out" "$err"
  echo "PASS"
  return 0
}

run_step() {
  local step_num="$1"
  local total
  local step
  local step_id
  local action_json
  local verification
  local rollback
  local expected
  local before_file
  local after_file
  local delta_file
  local diff_file
  local cmd_file
  local changed

  total="$(state_get '.plan_steps | length')"
  if (( step_num < 1 || step_num > total )); then
    echo "[step] invalid step index: $step_num"
    return 1
  fi

  step="$(jq -c ".plan_steps[$((step_num - 1))]" "$STATE_FILE")"
  step_id="$(jq -r '.step_id // empty' <<<"$step")"
  [[ -z "$step_id" ]] && step_id="$step_num"
  verification="$(jq -r '.verification // empty' <<<"$step")"
  rollback="$(jq -r '.rollback // empty' <<<"$step")"
  expected="$(jq -r '.expected_output // empty' <<<"$step")"
  action_json="$(jq -c '.action' <<<"$step")"

  if [[ -z "$action_json" || "$action_json" == "null" ]]; then
    echo "[step] empty action for step $step_num"
    append_question "step_${step_num}: empty action"
    state_set_string '.pending_action' 'await_next_approval'
    return 1
  fi

  echo "[step] step_id=$step_id"
  echo "[step] action=$action_json"
  echo "[step] expected=$expected"
  echo "[step] rollback=$rollback"

  before_file=$(mktemp)
  after_file=$(mktemp)
  delta_file=$(mktemp)
  diff_file=$(mktemp)
  cmd_file=$(mktemp)
  trap 'rm -f "$before_file" "$after_file" "$delta_file" "$diff_file" "$cmd_file"' RETURN
  collect_changed_files > "$before_file"

  if ! collect_step_actions "$action_json" > "$cmd_file"; then
    echo "[step] invalid action payload for step $step_num"
    append_question "step_${step_num}_invalid_action_payload"
    state_set_string '.pending_action' 'await_next_approval'
    state_set_expr '.plan_status = "approved"'
    state_set_expr '.updated_at = "'"$(now_ts)"'"'
    state_set_last_verification "$verification" "not_run"
    append_history "$step_id" "failure" "invalid action payload" "not_run" "revise"
    echo "[step] FAIL"
    return 1
  fi

  if [[ ! -s "$cmd_file" ]]; then
    echo "[step] no executable commands resolved for step $step_num"
    append_question "step_${step_num}_empty_action_commands"
    state_set_string '.pending_action' 'await_next_approval'
    state_set_expr '.plan_status = "approved"'
    state_set_expr '.updated_at = "'"$(now_ts)"'"'
    state_set_last_verification "$verification" "not_run"
    append_history "$step_id" "failure" "empty action command list" "not_run" "revise"
    echo "[step] FAIL"
    return 1
  fi

  local attempt=1
  local step_status="failure"
  while (( attempt <= 2 )); do
    local cmd_rc=0
    echo "[step] attempt $attempt/2"
    while IFS= read -r cmd; do
      [[ -z "$(printf '%s' "$cmd")" ]] && continue
      if ! run_command "$cmd" 1; then
        cmd_rc=1
        break
      fi
    done < "$cmd_file"

    if [[ $cmd_rc -eq 0 ]]; then
      step_status="success"
      break
    fi

    if (( attempt == 2 )); then
      break
    fi
    echo "[step] retrying"
    ((attempt++))
  done

  collect_changed_files > "$after_file"
  comm -3 "$before_file" "$after_file" > "$delta_file" || true
  changed="$(count_lines "$delta_file")"

  if [[ "$changed" -gt 2 ]]; then
    echo "[step] BLOCK: changed files count for this step is $changed, max is 2"
    step_status="failure"
  fi

  if [[ -d "$ROOT_DIR/.git" ]]; then
    git -C "$ROOT_DIR" diff -- . > "$diff_file"
  fi

  echo "[step] unified diff:"
  if [[ -s "$diff_file" ]]; then
    cat "$diff_file"
  else
    echo "(no tracked diff for this step)"
  fi

  local verify_result="not_run"
  if [[ "$step_status" == "success" ]]; then
    verify_result="$(run_verify_if_needed "$verification")"
    if [[ "$verify_result" != "PASS" && "$verify_result" != "not_run" ]]; then
      step_status="failure"
    fi
  fi

  if [[ "$step_status" == "success" ]]; then
    state_append '.completed_steps' "$step_id"
    state_set_number '.current_step' "$step_num"
    state_set_number '.next_step' "$((step_num + 1))"
    state_set_last_verification "$verification" "$verify_result"

    if (( step_num >= total )); then
      state_set_string '.plan_status' 'completed'
      state_set_string '.pending_action' 'completed'
    else
      state_set_string '.plan_status' 'approved'
      state_set_string '.pending_action' 'await_next_approval'
    fi

    state_set_expr '.open_questions = []'
    state_set_expr '.updated_at = "'"$(now_ts)"'"'
    append_history "$step_id" "success" "step complete" "$verify_result" "await_next"
    echo "[step] PASS"
    return 0
  fi

  append_question "step_${step_num}_failed: command failed or verify failed"
  state_set_string '.pending_action' 'await_next_approval'
  state_set_expr '.plan_status = "approved"'
  state_set_expr '.updated_at = "'"$(now_ts)"'"'
  state_set_last_verification "$verification" "$verify_result"
  append_history "$step_id" "failure" "step failed" "$verify_result" "revise"
  echo "[step] FAIL"
  return 1
}

append_question() {
  local msg="$1"
  state_append '.open_questions' "$msg"
}

ensure_open_json
ensure_plan_status() {
  local ps
  ps="$(state_get '.plan_status')"
  if [[ "$ps" == "null" || -z "$ps" ]]; then
    state_set_string '.plan_status' 'not_started'
  fi

  local pa
  pa="$(state_get '.pending_action')"
  if [[ "$pa" == "null" || -z "$pa" ]]; then
    state_set_string '.pending_action' 'need_plan_approval'
  fi
}

ensure_plan_status

COMMAND="${1:-status}"
ARG="${2:-}"

case "${COMMAND,,}" in
  status)
    status
    ;;
  approve_plan)
    ensure_plan_exists
    state_set_string '.plan_status' 'approved'
    state_set_string '.pending_action' 'step_ready'
    state_set_number '.next_step' 1
    state_set_expr '.open_questions = []'
    state_mark_updated
    echo "[plan] APPROVE_PLAN accepted"
    ;;
  revise_plan)
    ensure_plan_exists
    append_question "${ARG:-plan requires revision}"
    state_set_string '.plan_status' 'revise_requested'
    state_set_string '.pending_action' 'revise_requested'
    state_mark_updated
    echo "[plan] REVISE_PLAN"
    ;;
  cancel_task)
    append_question "${ARG:-user cancelled}"
    state_set_string '.plan_status' 'cancelled'
    state_set_string '.pending_action' 'cancelled'
    state_mark_updated
    echo "[plan] CANCEL_TASK"
    ;;
  approve_next)
    if [[ "$(state_get '.pending_action')" != "await_next_approval" ]]; then
      echo "[step] no pending step for approve-next"
      status
      exit 1
    fi
    state_set_string '.pending_action' 'step_ready'
    state_mark_updated
    echo "[step] APPROVE_NEXT accepted"
    ;;
  step)
    ensure_plan_exists
    if [[ "$(state_get '.plan_status')" != "approved" ]]; then
      echo "[step] plan is not approved"
      status
      exit 1
    fi
    if [[ "$(state_get '.pending_action')" != "step_ready" ]]; then
      echo "[step] waiting token mismatch: not in step_ready"
      status
      exit 1
    fi

    target_step="$ARG"
    if [[ -z "$target_step" ]]; then
      target_step="$(state_get '.next_step')"
    fi
    if ! [[ "$target_step" =~ ^[0-9]+$ ]]; then
      echo "[step] step must be integer"
      exit 1
    fi
    run_step "$target_step"
    status
    ;;
  *)
    if [[ "${COMMAND,,}" == "approve-plan" ]]; then
      exec "$0" APPROVE_PLAN "${ARG:-}"
    elif [[ "${COMMAND,,}" == "revise-plan" ]]; then
      exec "$0" REVISE_PLAN "${ARG:-}"
    elif [[ "${COMMAND,,}" == "cancel-task" ]]; then
      exec "$0" CANCEL_TASK "${ARG:-}"
    elif [[ "${COMMAND,,}" == "approve-next" ]]; then
      exec "$0" APPROVE_NEXT
    else
      cat <<'EOF_USAGE'
Usage:
  run_task.sh status
  run_task.sh APPROVE_PLAN
  run_task.sh REVISE_PLAN [note]
  run_task.sh CANCEL_TASK
  run_task.sh step [N]
  run_task.sh APPROVE_NEXT
EOF_USAGE
    fi
    ;;
esac