#!/bin/bash
# Wrapper script to run AI agents (Claude/Codex/Gemini/Bob/Grok) and output JSON statistics
#
# Model selection via environment variables:
#   CLAUDE_MODEL - Model for Claude (e.g., opus, sonnet, haiku)
#   CLAUDE_CODE_EFFORT_LEVEL - Effort level for Claude (low, medium, high, xhigh, max)
#   OPENAI_MODEL - Model for Codex/OpenAI (e.g., gpt-5.5)
#   GEMINI_MODEL - Model for Gemini (e.g., gemini-2.5-pro)
#   BOB_SHELL_MODEL - Model for Bob-Shell (optional; Bob handles model selection internally)
#   GROK_MODEL - Model for Grok Build (e.g., grok-4.3, grok-4.20-multi-agent-0309)
#   GROK_EFFORT_LEVEL - Effort level for Grok Build (low/medium/high for grok-4.3; low/medium/high/xhigh for multi-agent)
#   CODEX_REASONING_LEVEL - Reasoning level for Codex (low, medium, high, xhigh)
#   CODEX_MODEL_PROVIDER - Optional Codex model provider override (e.g., coderflow_managed_custom_endpoint)
#
# If model env vars are not set, agents use their default models.

# Source refreshed credentials (token refresh writes updated env vars here)
if [ -f /home/coder/.bash_env ]; then
  source /home/coder/.bash_env
fi

# First argument is the agent command, rest are passed to the agent
AGENT_CMD="$1"
shift

echo "[$(date +'%Y-%m-%d %H:%M:%S')] Starting agent: $AGENT_CMD" >&2

# Log model configuration for the active agent only
if [[ "$AGENT_CMD" == *"codex"* ]]; then
  if [[ -n "$OPENAI_MODEL" ]]; then
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Model: $OPENAI_MODEL" >&2
  fi
  if [[ -n "$CODEX_REASONING_LEVEL" ]]; then
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Reasoning level: $CODEX_REASONING_LEVEL" >&2
  fi
  if [[ -n "$CODEX_MODEL_PROVIDER" ]]; then
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Model provider: $CODEX_MODEL_PROVIDER" >&2
  fi
elif [[ "$AGENT_CMD" == *"gemini"* ]]; then
  if [[ -n "$GEMINI_MODEL" ]]; then
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Model: $GEMINI_MODEL" >&2
  fi
elif [[ "$AGENT_CMD" == *"bob"* ]]; then
  if [[ -n "$BOB_SHELL_MODEL" ]]; then
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Model: $BOB_SHELL_MODEL" >&2
  fi
elif [[ "$AGENT_CMD" == *"grok"* ]]; then
  if [[ -n "$GROK_MODEL" ]]; then
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Model: $GROK_MODEL" >&2
  fi
  if [[ -n "$GROK_EFFORT_LEVEL" ]]; then
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Effort level: $GROK_EFFORT_LEVEL" >&2
  fi
else
  if [[ -n "$CLAUDE_MODEL" ]]; then
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Model: $CLAUDE_MODEL" >&2
  fi
  if [[ -n "$CLAUDE_CODE_EFFORT_LEVEL" ]]; then
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Effort level: $CLAUDE_CODE_EFFORT_LEVEL" >&2
  fi
fi

START_TIME=$(date +%s%3N)

# Save stdin to temp file
TEMP_INPUT=$(mktemp)
cat > "$TEMP_INPUT"

mkdir -p /task-output

AGENT_PID_FILE="/task-output/.agent_pid"
AGENT_PID=""

cleanup_agent_state() {
    if [ -n "$AGENT_PID" ]; then
        rm -f "$AGENT_PID_FILE"
        AGENT_PID=""
    fi
}

wait_for_agent() {
    wait "$AGENT_PID"
    local exit_code=$?
    cleanup_agent_state
    return "$exit_code"
}

DEBUG_FIFO=""
STAMPER_PID=""

# Start a background filter that injects a wall-clock `timestamp` field into
# any timestamp-less JSON line emitted by the agent CLI before appending it
# to /task-output/debug-stream.jsonl. Without this, Claude Code and Codex
# both write events without timestamps, the server falls back to synthetic
# timestamps, and the per-turn elapsed badge is suppressed by design.
#
# A FIFO is used instead of a plain pipe so $AGENT_PID still points at the
# agent process (the rightmost command of the foreground pipeline at each
# callsite) and wait_for_agent keeps returning the agent's exit code.
start_debug_stream_stamper() {
    DEBUG_FIFO=$(mktemp -u /task-output/.debug-stream-fifo.XXXXXX)
    if ! mkfifo "$DEBUG_FIFO" 2>/dev/null; then
        echo "[$(date +'%Y-%m-%d %H:%M:%S')] Fatal: failed to create FIFO for debug-stream timestamping at $DEBUG_FIFO" >&2
        exit 1
    fi
    node /usr/local/bin/stamp-debug-stream.js < "$DEBUG_FIFO" >> /task-output/debug-stream.jsonl 2>&1 &
    STAMPER_PID=$!
}

# Wait for the stamper to drain the FIFO and exit. Must be called after the
# agent has exited (so the FIFO write end is closed and the stamper sees
# EOF) and before any code that reads /task-output/debug-stream.jsonl.
wait_debug_stream_stamper() {
    if [ -n "$STAMPER_PID" ]; then
        wait "$STAMPER_PID" 2>/dev/null
        STAMPER_PID=""
    fi
    if [ -n "$DEBUG_FIFO" ]; then
        rm -f "$DEBUG_FIFO"
        DEBUG_FIFO=""
    fi
}

cleanup_on_exit() {
    rm -f "$TEMP_INPUT"
    cleanup_agent_state
    wait_debug_stream_stamper
}

trap cleanup_on_exit EXIT

# UNIFIED APPROACH FOR ALL AGENTS:
# - Run agent once, stream output to debug-stream.jsonl
# - Capture exit code
# - Generate a simple success/failure result
# - Let the UI parse debug-stream.jsonl for detailed activity

start_debug_stream_stamper

# Codex requires exec subcommand for non-interactive use (or resume for continuing)
# Claude accepts stdin directly with --print --output-format stream-json
# Gemini uses -p flag for prompt with --yolo for auto-approval and --resume for continuing
if [[ "$AGENT_CMD" == *"codex"* ]]; then
    # Build model argument if OPENAI_MODEL is set
    MODEL_ARG=""
    if [[ -n "$OPENAI_MODEL" ]]; then
        MODEL_ARG="-m $OPENAI_MODEL"
    fi

    # Build reasoning level argument if CODEX_REASONING_LEVEL is set
    REASONING_ARG=""
    if [[ -n "$CODEX_REASONING_LEVEL" ]]; then
        REASONING_ARG="-c model_reasoning_effort=$CODEX_REASONING_LEVEL"
    fi

    # Build model provider override argument if CODEX_MODEL_PROVIDER is set
    PROVIDER_ARG=""
    if [[ -n "$CODEX_MODEL_PROVIDER" ]]; then
        PROVIDER_ARG="-c model_provider=$CODEX_MODEL_PROVIDER"
    fi

    # Pass instructions via stdin (using "-" prompt) to avoid ARG_MAX limits with large instructions.
    # Check if this is a resume operation (first arg is 'exec', second is 'resume')
    if [[ "$1" == "exec" ]] && [[ "$2" == "resume" ]]; then
        shift 2  # Remove 'exec' and 'resume' from $@
        cat "$TEMP_INPUT" | "$AGENT_CMD" $MODEL_ARG $REASONING_ARG $PROVIDER_ARG -a never -s danger-full-access -c approval_policy=never -c sandbox_mode=danger-full-access exec --skip-git-repo-check --json resume "$@" > "$DEBUG_FIFO" 2>&1 &
        AGENT_PID=$!
        echo "$AGENT_PID" > "$AGENT_PID_FILE"
        wait_for_agent
        EXIT_CODE=$?
    else
        cat "$TEMP_INPUT" | "$AGENT_CMD" $MODEL_ARG $REASONING_ARG $PROVIDER_ARG -a never -s danger-full-access exec --skip-git-repo-check --json "$@" - > "$DEBUG_FIFO" 2>&1 &
        AGENT_PID=$!
        echo "$AGENT_PID" > "$AGENT_PID_FILE"
        wait_for_agent
        EXIT_CODE=$?
    fi

elif [[ "$AGENT_CMD" == *"gemini"* ]]; then
    # Use GEMINI_MODEL if set, otherwise default to "pro"
    GEMINI_MODEL_TO_USE="${GEMINI_MODEL:-pro}"

    # Pass instructions via stdin to avoid ARG_MAX limits with large instructions.
    # Gemini's -p flag appends to stdin input, so we pipe stdin and omit the -p value.
    # Check if this is a resume operation (--resume flag in args)
    # Include both /task-output and /workspace directories so Gemini can access
    # judge data files and workspace repos regardless of working directory.
    # Headless Gemini runs also need an explicit workspace trust override.
    if [[ "$*" == *"--resume"* ]]; then
        cat "$TEMP_INPUT" | "$AGENT_CMD" --model "$GEMINI_MODEL_TO_USE" --resume latest --yolo --skip-trust --include-directories /task-output --include-directories /workspace --output-format stream-json > "$DEBUG_FIFO" 2>&1 &
        AGENT_PID=$!
        echo "$AGENT_PID" > "$AGENT_PID_FILE"
        wait_for_agent
        EXIT_CODE=$?
    else
        cat "$TEMP_INPUT" | "$AGENT_CMD" --model "$GEMINI_MODEL_TO_USE" --yolo --skip-trust --include-directories /task-output --include-directories /workspace --output-format stream-json "$@" > "$DEBUG_FIFO" 2>&1 &
        AGENT_PID=$!
        echo "$AGENT_PID" > "$AGENT_PID_FILE"
        wait_for_agent
        EXIT_CODE=$?
    fi

elif [[ "$AGENT_CMD" == *"bob"* ]]; then
    # Bob-Shell is a Gemini CLI fork - uses BOB_SHELL_MODEL env var
    # Only pass --model if explicitly configured; otherwise let Bob use its default
    BOB_MODEL_ARG=""
    if [[ -n "$BOB_SHELL_MODEL" ]]; then
        BOB_MODEL_ARG="--model $BOB_SHELL_MODEL"
    fi

    # Same pattern as Gemini: stdin for instructions, --yolo for auto-approval
    # Note: Bob's --resume mode does not read stdin for the follow-up prompt
    # (despite the help text suggesting it should). We must use -p flag instead.
    if [[ "$*" == *"--resume"* ]]; then
        "$AGENT_CMD" $BOB_MODEL_ARG --accept-license --resume latest --yolo --trust --include-directories /task-output --include-directories /workspace --output-format stream-json -p "$(cat "$TEMP_INPUT")" > "$DEBUG_FIFO" 2>&1 &
        AGENT_PID=$!
        echo "$AGENT_PID" > "$AGENT_PID_FILE"
        wait_for_agent
        EXIT_CODE=$?
    else
        cat "$TEMP_INPUT" | "$AGENT_CMD" $BOB_MODEL_ARG --accept-license --yolo --trust --include-directories /task-output --include-directories /workspace --output-format stream-json "$@" > "$DEBUG_FIFO" 2>&1 &
        AGENT_PID=$!
        echo "$AGENT_PID" > "$AGENT_PID_FILE"
        wait_for_agent
        EXIT_CODE=$?
    fi

elif [[ "$AGENT_CMD" == *"grok"* ]]; then
    # Grok Build CLI: reads the prompt from -p (or --single). The CLI does NOT
    # accept the prompt on stdin, so we pass the entire input as a single argument.
    # --output-format streaming-json emits newline-delimited events:
    #   {"type":"text","data":"..."}      response chunk
    #   {"type":"thought","data":"..."}   internal reasoning
    #   {"type":"end","stopReason":"...","sessionId":"...","requestId":"..."}
    # --yolo auto-approves all tool executions (required for unattended runs).
    GROK_MODEL_ARG=""
    if [[ -n "$GROK_MODEL" ]]; then
        GROK_MODEL_ARG="-m $GROK_MODEL"
    fi
    GROK_EFFORT_ARG=""
    if [[ -n "$GROK_EFFORT_LEVEL" ]]; then
        GROK_EFFORT_ARG="--effort $GROK_EFFORT_LEVEL"
    fi

    # GROK_TRANSPORT selects between ACP (Agent Client Protocol —
    # https://agentclientprotocol.com) and the older --output-format
    # streaming-json. ACP is the default because it surfaces tool_call
    # notifications with full diff/locations/kind that streaming-json
    # doesn't expose, giving the activity feed the per-file Edit/Bash/Read
    # items the other agents already have. Opt out with
    # GROK_TRANSPORT=streaming-json if ACP misbehaves.
    GROK_TRANSPORT_MODE="${GROK_TRANSPORT:-acp}"

    if [[ "$GROK_TRANSPORT_MODE" == "acp" ]]; then
        # ACP transport. Use `grok agent stdio` (newline-delimited JSON-RPC)
        # behind acp-client.js, which drives the prompt-turn lifecycle and
        # echoes every inbound frame to debug-stream.jsonl.
        # --reasoning-effort on `grok agent` is the ACP equivalent of --effort.
        GROK_REASONING_ARG=""
        if [[ -n "$GROK_EFFORT_LEVEL" ]]; then
            GROK_REASONING_ARG="--reasoning-effort $GROK_EFFORT_LEVEL"
        fi
        # Resume: pull the previous sessionId from debug-stream.jsonl, same
        # pattern as Claude --continue. Grok nests sessionId in different
        # places depending on event type (params.sessionId for notifications,
        # result.sessionId for session/new, result._meta.sessionId for the
        # session/prompt response in recent CLI versions), so recurse to find
        # it wherever it lives.
        ACP_RESUME_ARG=""
        if [[ "$*" == *"--resume"* ]] && [ -f /task-output/debug-stream.jsonl ]; then
            ACP_PREV_SID=$(grep -aE '"sessionId"' /task-output/debug-stream.jsonl | tail -1 | jq -r 'first(.. | objects | .sessionId? | select(type == "string" and . != "")) // empty' 2>/dev/null)
            if [ -n "$ACP_PREV_SID" ]; then
                ACP_RESUME_ARG="--resume $ACP_PREV_SID"
                echo "[$(date +'%Y-%m-%d %H:%M:%S')] ACP resuming session: $ACP_PREV_SID" >&2
            else
                echo "[$(date +'%Y-%m-%d %H:%M:%S')] ACP --resume requested but no prior sessionId found; starting fresh" >&2
            fi
        fi
        cat "$TEMP_INPUT" | node /usr/local/bin/acp-client.js $ACP_RESUME_ARG --cwd /workspace -- \
            "$AGENT_CMD" agent --always-approve $GROK_MODEL_ARG $GROK_REASONING_ARG stdio > "$DEBUG_FIFO" 2>&1 &
        AGENT_PID=$!
        echo "$AGENT_PID" > "$AGENT_PID_FILE"
        wait_for_agent
        EXIT_CODE=$?
    elif [[ "$*" == *"--resume"* ]]; then
        # Resume previous session if requested. Grok's --resume takes an optional
        # session ID; with no ID it resumes the most recent for the cwd. The
        # framework passes --resume as a flag, so we use the latter form.
        "$AGENT_CMD" $GROK_MODEL_ARG $GROK_EFFORT_ARG --yolo --cwd /workspace --output-format streaming-json --resume -p "$(cat "$TEMP_INPUT")" > "$DEBUG_FIFO" 2>&1 &
        AGENT_PID=$!
        echo "$AGENT_PID" > "$AGENT_PID_FILE"
        wait_for_agent
        EXIT_CODE=$?
    else
        "$AGENT_CMD" $GROK_MODEL_ARG $GROK_EFFORT_ARG --yolo --cwd /workspace --output-format streaming-json -p "$(cat "$TEMP_INPUT")" > "$DEBUG_FIFO" 2>&1 &
        AGENT_PID=$!
        echo "$AGENT_PID" > "$AGENT_PID_FILE"
        wait_for_agent
        EXIT_CODE=$?
    fi

else
    # Claude - pass model / effort via CLI flags when configured
    MODEL_ARGS=()
    if [[ -n "$CLAUDE_MODEL" ]]; then
        MODEL_ARGS=("--model" "$CLAUDE_MODEL")
    fi
    EFFORT_ARGS=()
    if [[ -n "$CLAUDE_CODE_EFFORT_LEVEL" ]]; then
        EFFORT_ARGS=("--effort" "$CLAUDE_CODE_EFFORT_LEVEL")
    fi

    # Resolve --continue to --resume <session_id> for reliable session continuation.
    # Claude Code's --continue flag may not find the previous session in --print mode,
    # so we explicitly extract the session_id from the last completed run.
    CLAUDE_ARGS=()
    for arg in "$@"; do
        if [[ "$arg" == "--continue" ]] && [ -f /task-output/debug-stream.jsonl ]; then
            SESSION_ID=$(grep '"type":"result"' /task-output/debug-stream.jsonl | tail -1 | jq -r '.session_id // empty' 2>/dev/null)
            if [ -n "$SESSION_ID" ]; then
                CLAUDE_ARGS+=("--resume" "$SESSION_ID")
                echo "[$(date +'%Y-%m-%d %H:%M:%S')] Resuming session: $SESSION_ID" >&2
            else
                CLAUDE_ARGS+=("--continue")
                echo "[$(date +'%Y-%m-%d %H:%M:%S')] No session_id found, falling back to --continue" >&2
            fi
        else
            CLAUDE_ARGS+=("$arg")
        fi
    done

    cat "$TEMP_INPUT" | "$AGENT_CMD" --print --output-format stream-json --verbose "${MODEL_ARGS[@]}" "${EFFORT_ARGS[@]}" "${CLAUDE_ARGS[@]}" > "$DEBUG_FIFO" 2>&1 &
    AGENT_PID=$!
    echo "$AGENT_PID" > "$AGENT_PID_FILE"
    wait_for_agent
    EXIT_CODE=$?
fi

# Drain the timestamping filter before reading debug-stream.jsonl below.
# The agent has exited, which closes the FIFO write end, which causes the
# stamper to see EOF and flush its remaining buffered output.
wait_debug_stream_stamper

END_TIME=$(date +%s%3N)
DURATION=$((END_TIME - START_TIME))

# Extract final response from stream for activity feed
# This ensures the conversation shows the agent's reply
TIMESTAMP=$(date -u +"%Y-%m-%dT%H:%M:%S.%3NZ")

if [ -f /task-output/debug-stream.jsonl ] && [ $EXIT_CODE -eq 0 ]; then
    # Try to extract final response based on agent type
    if [[ "$AGENT_CMD" == *"claude"* ]]; then
        # Claude: Look for last text block in assistant message
        RESPONSE=$(grep -a '"type":"assistant"' /task-output/debug-stream.jsonl | tail -1 | jq -r '.message.content[] | select(.type=="text") | .text' 2>/dev/null | head -20 || echo "")
    elif [[ "$AGENT_CMD" == *"gemini"* ]] || [[ "$AGENT_CMD" == *"bob"* ]]; then
        # Gemini/Bob-Shell: Extract last assistant message content (the actual response)
        # Format: {"type":"message","role":"assistant","content":"...","delta":true}
        RESPONSE=$(grep -a '"type":"message"' /task-output/debug-stream.jsonl | grep -a '"role":"assistant"' | tail -1 | jq -r '.content' 2>/dev/null || echo "")
    elif [[ "$AGENT_CMD" == *"codex"* ]]; then
        # Codex: Look for last agent_message or reasoning item
        RESPONSE=$(grep -a '"type":"item.completed"' /task-output/debug-stream.jsonl | tail -1 | jq -r 'select(.item.type=="agent_message") | .item.text' 2>/dev/null || \
                   grep -a '"type":"item.completed"' /task-output/debug-stream.jsonl | tail -1 | jq -r 'select(.item.type=="reasoning") | .item.text' 2>/dev/null | head -20 || echo "")
    elif [[ "$AGENT_CMD" == *"grok"* ]]; then
        # Grok streaming-json emits {"type":"text","data":"chunk"} repeatedly for the
        # response, terminated by {"type":"end",...}. Concatenate all text chunks.
        RESPONSE=$(grep -a '"type":"text"' /task-output/debug-stream.jsonl | jq -r 'select(.type=="text") | .data' 2>/dev/null | tr -d '\n' | head -c 2000 || echo "")
    fi

    # Clean up response (limit length for summary)
    if [ -n "$RESPONSE" ] && [ "$RESPONSE" != "null" ]; then
        RESULT_TEXT="$RESPONSE"
    else
        RESULT_TEXT="Task completed successfully"
    fi
else
    if [ $EXIT_CODE -eq 130 ]; then
        RESULT_TEXT="Task interrupted"
    else
        RESULT_TEXT="Task failed with exit code $EXIT_CODE"
    fi
fi

# Estimate tokens from debug stream size
if [ -f /task-output/debug-stream.jsonl ]; then
    TOKEN_EST=$(wc -c < /task-output/debug-stream.jsonl | awk '{print int($1/4)}')
else
    TOKEN_EST=0
fi

# Escape result text for JSON
ESCAPED_RESULT=$(echo "$RESULT_TEXT" | jq -Rs . 2>/dev/null || echo '"Task completed"')

# Output result JSON with actual agent response
# The server and UI will also read debug-stream.jsonl for detailed events
cat <<EOF
{"type":"result","result":$ESCAPED_RESULT,"duration_ms":${DURATION},"num_turns":1,"usage":{"output_tokens":${TOKEN_EST},"cache_read_input_tokens":0},"timestamp":"${TIMESTAMP}","exit_code":${EXIT_CODE}}
EOF

# Write exit code and finished time marker files for server monitoring
echo "$EXIT_CODE" > /task-output/.exit_code
date -Iseconds > /task-output/.finished_at

exit $EXIT_CODE
