#!/bin/bash
# Wrapper script to run AI agents (Claude/Codex/Gemini/Bob) 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, 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)
#   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
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"
}

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

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

# 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 "$@" >> /task-output/debug-stream.jsonl 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 "$@" - >> /task-output/debug-stream.jsonl 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 >> /task-output/debug-stream.jsonl 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 "$@" >> /task-output/debug-stream.jsonl 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")" >> /task-output/debug-stream.jsonl 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 "$@" >> /task-output/debug-stream.jsonl 2>&1 &
        AGENT_PID=$!
        echo "$AGENT_PID" > "$AGENT_PID_FILE"
        wait_for_agent
        EXIT_CODE=$?
    fi

else
    # Claude - pass model via --model flag if CLAUDE_MODEL is set
    MODEL_ARG=""
    if [[ -n "$CLAUDE_MODEL" ]]; then
        MODEL_ARG="--model $CLAUDE_MODEL"
    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_ARG "${CLAUDE_ARGS[@]}" >> /task-output/debug-stream.jsonl 2>&1 &
    AGENT_PID=$!
    echo "$AGENT_PID" > "$AGENT_PID_FILE"
    wait_for_agent
    EXIT_CODE=$?
fi

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 "")
    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
