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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
# shellcheck source=/dev/null
source "${ROOT_DIR}/src/app_v2.sh"

usage() {
  cat <<USAGE
usage: $0 [--json] [--require-evidence] [--max-steps N] [--state-file PATH] [--auto-start-daemon] ROOTS_CSV QUESTION [GOAL_PREDICATES_CSV]

Minimal agent loop v1:
  1) retrieve context
  2) ask model for next action (answer/refine_query/stop)
  3) optionally refine roots/goals and repeat
  4) produce grounded answer and persist run checkpoint

Options:
  --json                 print full final JSON payload
  --require-evidence     fail if no supporting edges are found
  --max-steps N          planner iterations before answer fallback (default: 3)
  --state-file PATH      write run state to this path
  --auto-start-daemon    start daemon if startup checks fail
USAGE
}

require_jq() {
  if ! command -v jq >/dev/null 2>&1; then
    echo "ai_agent_loop.sh: jq is required" >&2
    exit 2
  fi
}

ensure_daemon_ready() {
  if app_startup_checks >/dev/null 2>&1; then
    return 0
  fi

  if [[ "${auto_start_daemon}" == "1" ]]; then
    local daemon_backend="${AI_DAEMON_STORE_BACKEND:-fs}"
    local daemon_root="${AI_DAEMON_STORE_ROOT:-/tmp/amduat-asl-ai-agent}"
    local daemon_log="${AI_DAEMON_LOG_PATH:-/tmp/ai-agent-daemon.log}"
    echo "daemon not reachable; attempting startup via scripts/dev_start_daemon.sh" >&2
    STORE_BACKEND="${daemon_backend}" STORE_ROOT="${daemon_root}" SOCK="${SOCK}" SPACE="${SPACE}" \
      nohup "${ROOT_DIR}/scripts/dev_start_daemon.sh" >"${daemon_log}" 2>&1 &
    local daemon_boot_pid="$!"
    disown "${daemon_boot_pid}" 2>/dev/null || true
    local i
    for i in $(seq 1 80); do
      if app_startup_checks >/dev/null 2>&1; then
        return 0
      fi
      sleep 0.1
    done
    app_startup_checks >/dev/null 2>&1 || {
      echo "ai_agent_loop.sh: daemon still unreachable after startup attempt" >&2
      echo "see ${daemon_log} for startup logs" >&2
      return 1
    }
    return 0
  fi

  echo "ai_agent_loop.sh: daemon unreachable on SOCK=${SOCK}" >&2
  echo "hint: run ./scripts/dev_start_daemon.sh or pass --auto-start-daemon" >&2
  return 1
}

append_step() {
  local step_json="$1"
  steps_json="$(jq -c --argjson step "${step_json}" '. + [$step]' <<<"${steps_json}")"
}

extract_plan_json() {
  local model_out="$1"
  local raw_plan
  raw_plan="$(jq -r '.response // ""' <<<"${model_out}")"
  local normalized_plan
  normalized_plan="$(printf '%s\n' "${raw_plan}" \
    | sed -e '1s/^```[[:alnum:]_-]*[[:space:]]*$//' -e '$s/^```[[:space:]]*$//')"
  local parsed_plan
  parsed_plan="$(printf '%s' "${normalized_plan}" | jq -c '
    if type == "object" then .
    else {"action":"answer","reason":"planner_non_object"}
    end
  ' 2>/dev/null || printf '%s' '{"action":"answer","reason":"planner_parse_error"}')"
  jq -c '
    def clean_csv(v): (v // "" | tostring | gsub("[\\r\\n\\t]+";" ") | gsub(" +";" ") | sub("^ ";"") | sub(" $";""));
    . as $r
    | {
        action: (
          ($r.action // "answer" | tostring) as $a
          | if ($a == "answer" or $a == "refine_query" or $a == "stop") then $a else "answer" end
        ),
        next_roots_csv: clean_csv($r.next_roots_csv // ""),
        next_goals_csv: clean_csv($r.next_goals_csv // ""),
        reason: clean_csv($r.reason // "")
      }
  ' <<<"${parsed_plan}"
}

plan_next_action() {
  local question="$1"
  local roots_csv="$2"
  local goals_csv="$3"
  local retrieve_json="$4"
  local step_no="$5"

  local context_stats
  context_stats="$(jq -c '{nodes:(.nodes // [] | length), edges:(.edges // [] | length)}' <<<"${retrieve_json}")"

  local prompt
  prompt="$(cat <<PROMPT
You are an execution planner for a graph-grounded QA agent.
Decide the NEXT action only.

Rules:
- Return STRICT JSON object only.
- action must be one of: "answer", "refine_query", "stop".
- Use "refine_query" only if retrieval context is clearly insufficient.
- Keep next_roots_csv / next_goals_csv empty unless refining.
- Do not include markdown or prose outside JSON.

Current step: ${step_no}
Question: ${question}
Current roots_csv: ${roots_csv}
Current goals_csv: ${goals_csv}
Context stats: ${context_stats}

Required JSON schema:
{"action":"answer|refine_query|stop","next_roots_csv":"","next_goals_csv":"","reason":""}
PROMPT
)"

  local plan_model_out
  plan_model_out="$(app_ai_generate_json "${prompt}")"
  extract_plan_json "${plan_model_out}"
}

output_mode="text"
require_evidence=0
max_steps=3
auto_start_daemon=0
state_file=""

while [[ $# -gt 0 ]]; do
  case "$1" in
    --json)
      output_mode="json"
      shift
      ;;
    --require-evidence)
      require_evidence=1
      shift
      ;;
    --max-steps)
      [[ $# -ge 2 ]] || { usage >&2; exit 2; }
      max_steps="$2"
      shift 2
      ;;
    --state-file)
      [[ $# -ge 2 ]] || { usage >&2; exit 2; }
      state_file="$2"
      shift 2
      ;;
    --auto-start-daemon)
      auto_start_daemon=1
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    --)
      shift
      break
      ;;
    -*)
      usage >&2
      exit 2
      ;;
    *)
      break
      ;;
  esac
done

if [[ $# -lt 2 || $# -gt 3 ]]; then
  usage >&2
  exit 2
fi

roots_csv="$1"
question="$2"
goals_csv="${3:-}"

if ! [[ "${max_steps}" =~ ^[0-9]+$ ]] || [[ "${max_steps}" -lt 1 ]] || [[ "${max_steps}" -gt 8 ]]; then
  echo "ai_agent_loop.sh: --max-steps must be integer in [1,8]" >&2
  exit 2
fi

require_jq
app_init
ensure_daemon_ready

run_id="$(date +%Y%m%d-%H%M%S)-$$"
if [[ -z "${state_file}" ]]; then
  mkdir -p "${ROOT_DIR}/ai/runs"
  state_file="${ROOT_DIR}/ai/runs/agent-run-${run_id}.json"
fi

steps_json="[]"
final_answer_json=""
stop_reason="max_steps_reached"

step_no=1
while (( step_no <= max_steps )); do
  retrieve_out="$(app_retrieve_with_fallback "${roots_csv}" "${goals_csv}")" || {
    stop_reason="retrieve_failed"
    break
  }
  context_stats="$(jq -c '{nodes:(.nodes // [] | length), edges:(.edges // [] | length)}' <<<"${retrieve_out}")"

  plan_json="$(plan_next_action "${question}" "${roots_csv}" "${goals_csv}" "${retrieve_out}" "${step_no}")"
  plan_action="$(jq -r '.action' <<<"${plan_json}")"
  next_roots="$(jq -r '.next_roots_csv // ""' <<<"${plan_json}")"
  next_goals="$(jq -r '.next_goals_csv // ""' <<<"${plan_json}")"

  step_record="$(jq -nc \
    --argjson step "${step_no}" \
    --arg roots_csv "${roots_csv}" \
    --arg goals_csv "${goals_csv}" \
    --argjson context "${context_stats}" \
    --argjson plan "${plan_json}" \
    '{step:$step,roots_csv:$roots_csv,goals_csv:$goals_csv,context:$context,plan:$plan}')"
  append_step "${step_record}"

  if [[ "${plan_action}" == "refine_query" ]]; then
    if [[ -n "${next_roots}" ]]; then
      roots_csv="${next_roots}"
    fi
    if [[ -n "${next_goals}" ]]; then
      goals_csv="${next_goals}"
    fi
    step_no=$(( step_no + 1 ))
    continue
  fi

  if [[ "${plan_action}" == "stop" ]]; then
    stop_reason="planner_stop"
    break
  fi

  if final_answer_json="$(app_ai_answer_json "${roots_csv}" "${question}" "${goals_csv}" "${require_evidence}")"; then
    stop_reason="answered"
    break
  fi

  stop_reason="answer_failed"
  break
done

if [[ -z "${final_answer_json}" ]]; then
  final_answer_json="$(jq -nc --arg msg "Agent loop ended without answer (${stop_reason})." '{response:$msg,done_reason:"agent_stopped"}')"
fi

run_json="$(jq -nc \
  --arg run_id "${run_id}" \
  --arg question "${question}" \
  --arg initial_roots_csv "$1" \
  --arg initial_goals_csv "${3:-}" \
  --arg final_roots_csv "${roots_csv}" \
  --arg final_goals_csv "${goals_csv}" \
  --arg stop_reason "${stop_reason}" \
  --argjson require_evidence "$( [[ "${require_evidence}" == "1" ]] && echo true || echo false )" \
  --argjson steps "${steps_json}" \
  --argjson final_answer "${final_answer_json}" \
  '{
      run_id:$run_id,
      input:{
        question:$question,
        roots_csv:$initial_roots_csv,
        goals_csv:$initial_goals_csv,
        require_evidence:$require_evidence
      },
      final_query:{
        roots_csv:$final_roots_csv,
        goals_csv:$final_goals_csv
      },
      stop_reason:$stop_reason,
      steps:$steps,
      final_answer:$final_answer
    }')"

printf '%s\n' "${run_json}" > "${state_file}"

if [[ "${output_mode}" == "json" ]]; then
  printf '%s\n' "${run_json}"
else
  jq -r '.final_answer.response // "No response"' <<<"${run_json}"
  echo
  echo "state_file=${state_file}"
fi