aiwb

#!/usr/bin/env bash
# aiwb - AIworkbench: Multi-Model AI Orchestration Platform
# Version: 2.0.0
# Universal compatibility: Linux, macOS, Termux
# License: GPL-3.0

set -o pipefail

# ============================================================================
# CWD SAFETY (Termux/Android)
# ============================================================================
# When spawned from a bot or after an Android app switch the original CWD may
# no longer exist, causing pwd/getcwd to fail.  Fix it early.
if ! pwd >/dev/null 2>&1; then
    cd "${HOME}" 2>/dev/null || cd / 2>/dev/null || true
fi

# ============================================================================
# INTERRUPT HANDLING
# ============================================================================

# Global flag to track if we're in an API call
AIWB_IN_API_CALL=0

# Cleanup function for graceful exit
cleanup_on_interrupt() {
    echo ""
    echo ""
    echo "Interrupted! Cleaning up..."

    # Kill all background jobs (more reliable than pkill, especially in Termux)
    local bg_jobs=$(jobs -p 2>/dev/null)
    if [[ -n "$bg_jobs" ]]; then
        # First try graceful termination
        kill $bg_jobs 2>/dev/null || true
        # Give processes 0.5s to exit gracefully
        sleep 0.5
        # Force kill any remaining processes
        kill -9 $bg_jobs 2>/dev/null || true
    fi

    # Also try pkill as backup for any orphaned curl processes
    pkill -P $$ curl 2>/dev/null || true

    # Clean up temp files (respect $TMPDIR for Termux)
    rm -f "${TMPDIR:-/tmp}"/aiwb_curl_* 2>/dev/null || true

    echo "Goodbye!"
    exit 130  # Standard exit code for SIGINT
}

# Set up trap for Ctrl+C (SIGINT) and Ctrl+Z (SIGTSTP)
trap cleanup_on_interrupt INT TERM

# ============================================================================
# BOOTSTRAP & LIBRARY LOADING
# ============================================================================

# ============================================================================
# BOOTSTRAP & LIBRARIES
# ============================================================================

# 1. Resolve where this script is actually located
SOURCE="${BASH_SOURCE[0]}"
while [ -h "$SOURCE" ]; do 
    DIR="$( cd -P "$( dirname "$SOURCE" )" >/dev/null 2>&1 && pwd )"
    SOURCE="$(readlink "$SOURCE")"
    [[ $SOURCE != /* ]] && SOURCE="$DIR/$SOURCE" 
done
SCRIPT_DIR="$( cd -P "$( dirname "$SOURCE" )" >/dev/null 2>&1 && pwd )"

# 2. Define Library Directory relative to script
LIB_DIR="${SCRIPT_DIR}/lib"
AIWB_LIB_DIR="${AIWB_LIB_DIR:-$LIB_DIR}"

# 3. Source all libraries found (with deduplication)
# We check LIB_DIR first (standard install), then fallback to local
declare -A seen_paths
for path in "$LIB_DIR" "$AIWB_LIB_DIR" "$HOME/.local/bin/lib"; do
    # Resolve to absolute path and deduplicate
    if [[ -d "$path" ]]; then
        abs_path=$(cd "$path" 2>/dev/null && pwd)
        if [[ -n "$abs_path" ]] && [[ -z "${seen_paths[$abs_path]:-}" ]]; then
            seen_paths[$abs_path]=1
            for lib in "$abs_path"/*.sh; do
                if [[ -f "$lib" ]]; then
                    source "$lib"
                fi
            done
        fi
    fi
done
unset seen_paths abs_path

# Safety: clear AIWB_HEADLESS unless this invocation is actually 'aiwb headless'.
# Prevents a leaked AIWB_HEADLESS=1 in the caller's environment from silently
# disabling the TUI for interactive sessions.
if [[ "${1:-}" != "headless" ]]; then
    unset AIWB_HEADLESS
fi



# ============================================================================
# INITIALIZATION
# ============================================================================

# Check dependencies
check_dependencies() {
    local missing=()

    for cmd in bash jq curl; do
        have "$cmd" || missing+=("$cmd")
    done

    if [[ ${#missing[@]} -gt 0 ]]; then
        err "Missing required dependencies: ${missing[*]}"
        echo ""
        echo "Install with:"
        if is_termux; then
            echo "  pkg install ${missing[*]}"
        elif is_macos; then
            echo "  brew install ${missing[*]}"
        else
            echo "  sudo apt install ${missing[*]}  # Debian/Ubuntu"
            echo "  sudo pacman -S ${missing[*]}    # Arch"
            echo "  sudo dnf install ${missing[*]}  # Fedora"
        fi
        exit "$E_DEPENDENCY"
    fi

    # Optional but recommended
    if ! have gum; then
        debug "gum not installed - using fallback UI"
    fi
}

# Initialize workspace and config
initialize() {
    debug "Initializing AIWB..."

    check_dependencies
    debug "After check_dependencies..."

    init_workspace
    debug "After init_workspace..."

    load_config
    debug "After load_config..."

    load_session
    debug "After load_session..."

    debug "About to load keys..."
    load_keys || true
    debug "After load keys..."

    # Auto-fix CRLF if on Termux
    if is_termux && [[ -d "$HOME/.local/bin" ]]; then
        fix_crlf_scripts
    fi

    # Skip automatic git detection to avoid delays in Termux
    # Users can manually scan repos with /scanrepo or /smartscan commands
    AIWB_REPO_ENABLED=false
    debug "Skipped detect_repo for faster startup"
}

# Fix CRLF line endings (Termux compatibility)
fix_crlf_scripts() {
    local bin_dir="$HOME/.local/bin"
    local fixed=0

    for script in "$bin_dir"/*.sh; do
        [[ -f "$script" ]] || continue
        if grep -q $'\r' "$script" 2>/dev/null; then
            sed -i 's/\r$//' "$script"
            ((fixed++))
        fi
    done

    if ((fixed > 0)); then
        debug "Fixed CRLF in $fixed scripts"
        hash -r
    fi
}

# ============================================================================
# COMMANDS
# ============================================================================

# Show help
cmd_help() {
    ui_header "AIWB - AI Workbench v2.0"

    cat <<'EOF'
A multi-model AI orchestration platform for the command line.

USAGE:
  aiwb [command] [options]

MODE-BASED WORKFLOWS (Recommended! 🎯):
  Start interactive session, then use modes:
    aiwb                    # Start interactive
    > /make                 # Generate from scratch
    > /tweak                # Modify existing code
    > /debug                # Find and fix bugs

  Inside each mode:
    prompt <text>           Set text instructions
    instruct <file.md>      Load from file
    model                   Choose AI model
    uploads <files/dirs>    Add context
    check [instructions]    Configure verification
    status                  Show mode status
    run                     Execute with cost dialog

QUICK COMMANDS:
  quick <description>     One-shot generation with auto-verify
  wizard                  Interactive guided workflow
  improve <suggestion>    Improve last output

STANDARD COMMANDS:
  (no command)         Start interactive chat
  chat                 Interactive chat mode
  generate <task>      Generate for a task
  verify <task>        Verify/review content
  refine <task>        Multi-iteration loop

CONFIGURATION:
  keys                 Manage API keys
  settings             Configure provider/model
  status               Show overall status
  doctor               Check system health

REPOSITORY SCANNING:
  scanrepo             Deep scan entire repository for full context
  smartscan            Quick scan of key files (configs, docs, main files)

GITHUB (like Claude Code):
  github               Interactive GitHub menu
  github status        Show git status
  github clone <repo>  Clone a repository
  github commit <msg>  Commit changes
  github push          Push to remote
  github pull          Pull from remote
  github sync          Sync with remote (fetch, pull/push as needed)
  github pr create     Create pull request
  github issue create  Create issue

UTILITY:
  history              Session history
  costs                Cost breakdown
  templates            Browse templates
  security-audit       Security audit

OPTIONS:
  --provider <name>    Use specific provider (gemini/claude/openai/groq/xai)
  --model <name>       Use specific model
  --debug              Enable debug output
  --help, -h           Show this help
  --version, -v        Show version

EXAMPLES:
  # Repository analysis
  cd my-repo && aiwb scanrepo    # Deep scan
  cd my-repo && aiwb smartscan   # Quick overview

  # Mode-based (powerful and flexible!)
  aiwb
  > /make
  make> prompt "Create REST API for users"
  make> model              # Choose Claude
  make> uploads ./docs/ ./src/models/
  make> check "Focus on security"
  make> run

  # Quick one-liners
  aiwb quick "Create a password generator"
  aiwb wizard
  aiwb improve "Add error handling"

  # Setup
  aiwb keys
  aiwb settings

Get started: aiwb wizard (beginners) or aiwb + /make (power users)
Learn more: https://github.com/juanitto-maker/AIworkbench

EOF
}

# Main interactive chat interface
cmd_chat() {
    clear 2>/dev/null || true  # Don't fail if TERM not set
    ui_logo

    # Show status
    ui_show_status
    echo ""

    # Check if keys are set
    local provider
    provider="$(config_get model_provider)"

    if ! has_api_key "$provider"; then
        echo ""
        echo "⚠️  No API key found for $provider."
        echo "   Please set up your API keys to use AIWB."
        echo ""
        echo "   Run: aiwb keys"
        echo "   Or use /keys command in chat"
        echo ""
    fi

    # Check for existing context from previous session
    if type context_state_exists &>/dev/null && context_state_exists; then
        local created_at
        created_at=$(context_state_get "created_at" 2>/dev/null || echo "")

        if [[ -n "$created_at" ]]; then
            local created_epoch now_epoch age_days age_msg
            created_epoch=$(date -d "$created_at" +%s 2>/dev/null || date -j -f "%Y-%m-%dT%H:%M:%S" "${created_at:0:19}" +%s 2>/dev/null || echo "0")
            now_epoch=$(date +%s)
            age_days=$(( (now_epoch - created_epoch) / 86400 ))

            if [[ $age_days -eq 0 ]]; then
                age_msg="today"
            elif [[ $age_days -eq 1 ]]; then
                age_msg="yesterday"
            else
                age_msg="$age_days days ago"
            fi

            echo ""
            warn "Found context from previous session ($age_msg)"

            if ui_confirm "Resume with previous context?" "yes"; then
                # Get list of available context files
                local -a available_files
                mapfile -t available_files < <(context_state_list_files)

                if [[ ${#available_files[@]} -gt 0 ]]; then
                    echo ""
                    echo "Select which context files to load:"
                    echo ""

                    # Show file selection UI
                    local -a selected_files
                    mapfile -t selected_files < <(ui_select_context_files "${available_files[@]}")

                    # Load selected files
                    if [[ ${#selected_files[@]} -gt 0 ]]; then
                        context_state_load_into_mode "${selected_files[@]}"
                        echo ""
                        context_state_show
                    else
                        echo "No files selected, starting with empty context"
                    fi

                    # Suggest /context command (dimmed)
                    echo ""
                    echo -e "\033[2m💡 Tip: Use /context command to add more files later\033[0m"
                else
                    echo "No context files found in previous session"
                fi
            fi
            echo ""
        fi
    fi

    # Quick security check (silent if no issues)
    # Disabled automatic check on startup to avoid false positives
    # Users can run: aiwb security audit
    # if type audit_git_exposure &>/dev/null; then
    #     # Run git exposure check silently - only outputs if issues found
    #     audit_git_exposure 2>/dev/null || true
    # fi

    # Start chat loop
    chat_loop
}











# Rotate old log files to save storage
rotate_logs() {
    local logs_dir="$1"
    local max_logs="${2:-10}"  # Keep only last 10 logs by default
    local max_size="${3:-10485760}"  # 10MB max per log

    # Remove logs older than max count
    local logs=($(ls -t "$logs_dir"/chat_*.log 2>/dev/null))
    if [[ ${#logs[@]} -gt $max_logs ]]; then
        for ((i=max_logs; i<${#logs[@]}; i++)); do
            rm -f "${logs[$i]}" 2>/dev/null
        done
    fi

    # Truncate logs larger than max size
    for log in "$logs_dir"/chat_*.log; do
        [[ ! -f "$log" ]] && continue
        local size=$(stat -c%s "$log" 2>/dev/null || stat -f%z "$log" 2>/dev/null || echo 0)
        if [[ $size -gt $max_size ]]; then
            # Keep beginning and end for context (500 lines each)
            head -500 "$log" > "$log.tmp"
            echo "... [log truncated - $(date)] ..." >> "$log.tmp"
            tail -500 "$log" >> "$log.tmp"
            mv "$log.tmp" "$log"
        fi
    done
}

# Chat loop
chat_loop() {
    local workspace logs_dir
    workspace="$(get_workspace)"
    logs_dir="$workspace/logs"
    local chat_log="$logs_dir/chat_$(date +%Y%m%d_%H%M%S).log"

    # Rotate old logs to save space
    rotate_logs "$logs_dir" "$AIWB_LOG_RETENTION" "$AIWB_MAX_LOG_SIZE"

    msg "Chat started. Type /help for commands or /exit to quit."
    echo ""

    while true; do
        local input

        # Get input - Lazy Gum: terminal stays scrollable until user starts typing
        if [[ -t 0 ]] && [[ "${AIWB_TEST_MODE:-0}" != "1" ]] && $GUM_AVAILABLE; then
            # Show minimal prompt - terminal is free to scroll
            printf "\n${DIM}[Scroll freely, or press any key to type...]${RESET}\n"

            # Wait for first keystroke without locking terminal
            local first_char
            read -n1 -s first_char || continue

            # Clear the prompt line
            echo -ne "\r\033[K\033[1A\r\033[K"

            # Now activate gum with the first character pre-filled
            if [[ -n "$first_char" ]]; then
                input=$(gum input --value "$first_char" --placeholder "Message or /command..." --width 80) || continue
            else
                # If first char was just Enter, show gum empty
                input=$(gum input --placeholder "Message or /command..." --width 80) || continue
            fi
        else
            # Use safe_read for piped input or Termux compatibility
            printf "\033[0m> "  # Reset colors before prompt
            safe_read input || continue
        fi

        [[ -z "$input" ]] && continue

        # Log input
        echo "[$(date -Iseconds)] USER: $input" >> "$chat_log"

        # Display user input
        echo -e "\n${GREEN}${BOLD}User:${RESET} $input"

        # Handle commands
        if [[ "$input" == /* ]]; then
            handle_slash_command "$input"
            local cmd_exit=$?
            # If slash command was interrupted, exit chat loop
            if [[ $cmd_exit -eq 130 ]]; then
                cleanup_on_interrupt
            fi
        else
            # Send to AI with smart routing (auto-detects intent for chat-first workflow)
            # Fallback to basic chat if routed version fails
            local response_output
            if type handle_chat_message_routed &>/dev/null; then
                response_output=$(handle_chat_message_routed "$input") || response_output=$(handle_chat_message "$input")
            else
                response_output=$(handle_chat_message "$input")
            fi
            local chat_exit=$?

            # If interrupted, exit chat loop
            if [[ $chat_exit -eq 130 ]]; then
                cleanup_on_interrupt
            fi

            # Only process output if not interrupted
            if [[ $chat_exit -eq 0 ]]; then
                # Display already happened in handle_chat_message (to stderr)
                # response_output contains only the response text for logging

                # Log response with truncation for large responses
                local response_lines=$(echo "$response_output" | wc -l)
                if [[ $response_lines -gt $AIWB_RESPONSE_LINES_THRESHOLD ]]; then
                    # Log only first 50 and last 50 lines for very long responses
                    {
                        echo "$response_output" | head -50
                        echo "... [truncated $(($response_lines - 100)) lines] ..."
                        echo "$response_output" | tail -50
                    } >> "$chat_log"
                else
                    echo "$response_output" >> "$chat_log"
                fi
            fi
        fi

        # No gap - messages flow continuously for maximum screen usage
    done
}

# Handle slash commands
handle_slash_command() {
    local cmd="$1"
    shift || true

    # Extract command and arguments
    local base_cmd="${cmd%% *}"
    local args="${cmd#* }"
    [[ "$args" == "$base_cmd" ]] && args=""

    case "$base_cmd" in
        /help)
            show_chat_help
            ;;
        /make)
            # Check if this is the old quick command or new mode
            if [[ -n "$args" ]]; then
                # Old style: /make "description"
                cmd_make "$args"
            else
                # New style: enter make mode
                mode_loop "make"
            fi
            ;;
        /tweak)
            mode_loop "tweak"
            ;;
        /debug)
            mode_loop "debug"
            ;;
        /keys)
            cmd_keys
            ;;
        /models)
            cmd_models
            ;;
        /status)
            cmd_status_overview
            ;;
        /context)
            cmd_context "$args"
            ;;
        /estimate)
            cmd_estimate
            ;;
        /generate)
            cmd_generate
            ;;
        /verify)
            cmd_verify "$args"
            ;;
        /refine)
            cmd_refine
            ;;
        /history)
            cmd_history
            ;;
        /costs)
            cmd_costs
            ;;
        /costreset)
            cmd_costreset
            ;;
        /github|/gh)
            cmd_github "$args"
            ;;
        /edit)
            cmd_edit "$args"
            ;;
        /repo)
            cmd_repo "$args"
            ;;
        /scanrepo)
            cmd_scanrepo
            ;;
        /smartscan)
            cmd_smartscan
            ;;
        /contextload)
            cmd_contextload
            ;;
        /contextsave)
            cmd_contextsave
            ;;
        /contextshow)
            cmd_contextshow
            ;;
        /contextclear)
            cmd_contextclear
            ;;
        /contextrefresh)
            cmd_contextrefresh
            ;;
        /contextremove)
            cmd_contextremove
            ;;
        /quick)
            cmd_quick "$args"
            ;;
        /wizard)
            cmd_wizard
            ;;
        /improve)
            cmd_improve "$args"
            ;;
        /templates)
            cmd_templates
            ;;
        /swarm)
            cmd_swarm "$args"
            ;;
        /clear)
            clear 2>/dev/null || true
            ui_logo
            ;;
        /exit|/quit)
            # Require explicit confirmation to exit
            local confirm_exit
            if $GUM_AVAILABLE; then
                # Force user to make explicit choice - affirmative required
                if gum confirm "Are you sure you want to exit?" --affirmative="Yes, exit" --negative="No, stay"; then
                    msg "Goodbye!"
                    exit 0
                else
                    info "Exit cancelled"
                fi
            else
                # Fallback: require explicit 'yes' response
                confirm_exit=$(ui_input "Type 'yes' to exit, or press Enter to cancel:")
                if [[ "$confirm_exit" == "yes" ]]; then
                    msg "Goodbye!"
                    exit 0
                else
                    info "Exit cancelled"
                fi
            fi
            ;;
        *)
            warn "Unknown command: $base_cmd"
            echo "Type /help for available commands"
            ;;
    esac
}

# Show chat help
show_chat_help() {
    cat <<'EOF'
CODE EDITING (Like Claude Code!):
  /edit <file> <task>   Edit a specific file with AI
  /edit <task>          Smart edit (AI finds files to modify)
  /repo                 Show current repository info
  /repo none            Disable repo mode
  /scanrepo             Deep scan current folder (works in ANY directory!)
  /smartscan            Quick scan (configs, docs, main files only)

  Just type what you want changed, and AI will edit your files directly!
  Example: "Add dark mode to the settings page"

CONTEXT MANAGEMENT:
  /contextshow          Show current context state
  /contextload          Load saved context into session
  /contextsave          Save current context to state
  /contextclear         Clear all context and start fresh
  /contextrefresh       Re-scan repository and update context
  /contextremove        Remove specific file from context

MODE-BASED WORKFLOWS:
  /make       Generate code from scratch (enter mode)
  /tweak      Modify/update existing code (enter mode)
  /debug      Find and fix errors (enter mode)

QUICK COMMANDS:
  /quick <description>    One-shot generation with auto-verify
  /wizard                 Interactive guided workflow
  /improve <suggestion>   Improve last output
  /templates              Browse and use templates

GITHUB:
  /github            Interactive GitHub menu
  /github status     Show git status
  /github commit     Commit changes
  /github push       Push to remote
  /github pull       Pull from remote
  /github sync       Sync with remote (fetch, pull/push as needed)
  /github pr create  Create pull request

MANAGEMENT:
  /status     Show overall status
  /history    Show chat history
  /costs      Show cost breakdown
  /costreset  Reset cost counter to $0.00
  /swarm      Configure swarm mode (multi-agent processing)
              Interactive menu to:
              - Toggle swarm on/off
              - Select strategy (auto/mapreduce/hierarchical)
              - Configure worker models
              - Configure aggregator models
              - Set worker count

CONFIGURATION:
  /keys       Manage API keys
  /models     Configure provider/model
  /clear      Clear screen
  /help       Show this help
  /exit       Quit

WORKFLOW:
  1. cd into your repo
  2. Run: aiwb
  3. Ask for changes: "Add feature X to file Y"
  4. Review diff and confirm
  5. /github commit "message"
  6. /github push
EOF
}

# ============================================================================
# STATUS FOOTER HELPERS
# ============================================================================

# Get context summary for status display
get_context_summary() {
    # Show ACTUAL context being used (MODE_UPLOADS), not just what's in .context_state
    # This ensures the footer accurately reflects what the API will receive
    if [[ ${#MODE_UPLOADS[@]} -gt 0 ]]; then
        local file_count=${#MODE_UPLOADS[@]}
        local image_count=0

        # Count images if get_context_images is available
        if type get_context_images &>/dev/null; then
            while IFS= read -r img; do
                [[ -n "$img" ]] && ((image_count++))
            done < <(get_context_images)
        fi

        if [[ $image_count -gt 0 ]]; then
            echo "$file_count file(s) ($image_count image(s))"
        else
            echo "$file_count file(s)"
        fi
    else
        # No in-memory context loaded, but check if persistent context exists (not loaded)
        if type context_state_exists &>/dev/null && context_state_exists; then
            local scan_count
            scan_count=$(context_state_get "last_scan.file_count" 2>/dev/null)
            if [[ -n "$scan_count" && "$scan_count" != "null" && "$scan_count" != "0" ]]; then
                # Show that context exists but is not loaded
                echo "none (saved: $scan_count files - use /contextload)"
            else
                echo "none"
            fi
        else
            echo "none"
        fi
    fi
}

# Get total session cost
get_session_cost() {
    local workspace logs_dir usage_log
    workspace="$(get_workspace 2>/dev/null)" || workspace="$HOME/.aiwb/workspace"
    logs_dir="$workspace/logs"
    usage_log="$logs_dir/usage.jsonl"

    if [[ ! -f "$usage_log" ]]; then
        echo "0.00"
        return
    fi

    # Sum costs from current session (today's entries)
    local total_cost
    if command -v jq &>/dev/null; then
        total_cost=$(jq -s 'map(.cost) | add // 0' "$usage_log" 2>/dev/null)
        # Ensure we have a valid number, default to 0 if empty or invalid
        if [[ -z "$total_cost" ]] || [[ ! "$total_cost" =~ ^[0-9]*\.?[0-9]+$ ]]; then
            total_cost="0"
        fi
        # If cost is > 0 but < 0.01, show minimum of 0.01
        if command -v bc &>/dev/null; then
            if (( $(echo "$total_cost > 0" | bc -l 2>/dev/null) )) && (( $(echo "$total_cost < 0.01" | bc -l 2>/dev/null) )); then
                echo "0.01"
            else
                printf "%.2f" "$total_cost" 2>/dev/null || echo "0.00"
            fi
        else
            printf "%.2f" "$total_cost" 2>/dev/null || echo "0.00"
        fi
    else
        echo "0.00"
    fi
}

# Get message count for current session
get_message_count() {
    local workspace logs_dir
    workspace="$(get_workspace 2>/dev/null)" || workspace="$HOME/.aiwb/workspace"
    logs_dir="$workspace/logs"

    # Count "USER:" lines in most recent chat log
    local latest_log
    latest_log=$(ls -t "$logs_dir"/chat_*.log 2>/dev/null | head -1)

    if [[ -f "$latest_log" ]]; then
        grep -c "^\\[.*\\] USER:" "$latest_log" 2>/dev/null || echo "0"
    else
        echo "0"
    fi
}

# Get last response cost
get_last_response_cost() {
    local workspace logs_dir
    workspace="$(get_workspace 2>/dev/null)" || workspace="$HOME/.aiwb/workspace"
    logs_dir="$workspace/logs"
    local usage_log="$logs_dir/usage.jsonl"

    if [[ ! -f "$usage_log" ]]; then
        echo "0.00"
        return
    fi

    # Get the last entry's cost
    local last_cost
    if command -v jq &>/dev/null; then
        last_cost=$(tail -1 "$usage_log" | jq -r '.cost // 0' 2>/dev/null)
        # Ensure we have a valid number, default to 0 if empty or invalid
        if [[ -z "$last_cost" ]] || [[ ! "$last_cost" =~ ^[0-9]*\.?[0-9]+$ ]]; then
            last_cost="0"
        fi
        # If cost is > 0 but < 0.01, show minimum of 0.01
        if command -v bc &>/dev/null; then
            if (( $(echo "$last_cost > 0" | bc -l 2>/dev/null) )) && (( $(echo "$last_cost < 0.01" | bc -l 2>/dev/null) )); then
                echo "0.01"
            else
                printf "%.2f" "$last_cost" 2>/dev/null || echo "0.00"
            fi
        else
            printf "%.2f" "$last_cost" 2>/dev/null || echo "0.00"
        fi
    else
        echo "0.00"
    fi
}

# Show status footer after AI response
show_status_footer() {
    local provider model
    provider="$(config_get model_provider)"
    model="$(config_get model_name)"

    # Get terminal width (default 80 for narrow screens)
    local term_width
    term_width=$(tput cols 2>/dev/null || echo "80")

    # Build model display (provider/model)
    local model_display="${provider}/${model}"

    # Get context summary
    local context_summary
    context_summary=$(get_context_summary)

    # Get cost stats
    local session_cost message_count last_cost
    session_cost=$(get_session_cost)
    last_cost=$(get_last_response_cost)
    message_count=$(get_message_count)

    # Line 1: Model only
    local line1_plain="Model: ${model_display}"
    if [[ ${#line1_plain} -gt $term_width ]]; then
        local max_len=$((term_width - 3))
        local truncated="${line1_plain:0:$max_len}..."
        line1="${truncated}"
    else
        line1="Model: ${BOLD}${model_display}${RESET}"
    fi

    # Line 2: Context only
    local line2_plain="Context: ${context_summary}"
    if [[ ${#line2_plain} -gt $term_width ]]; then
        local max_len=$((term_width - 3))
        local truncated="${line2_plain:0:$max_len}..."
        line2="${truncated}"
    else
        line2="Context: ${DIM}${context_summary}${RESET}"
    fi

    # Line 3: Cost (last + session total) + messages
    local cost_color="${GREEN}"
    local session_color="${GREEN}"

    if command -v bc &>/dev/null; then
        # Color last response cost
        if (( $(echo "$last_cost > 0.10" | bc -l 2>/dev/null) )); then
            cost_color="${RED}"
        elif (( $(echo "$last_cost > 0.01" | bc -l 2>/dev/null) )); then
            cost_color="${YELLOW}"
        fi

        # Color session total
        if (( $(echo "$session_cost > 1.00" | bc -l 2>/dev/null) )); then
            session_color="${RED}"
        elif (( $(echo "$session_cost > 0.10" | bc -l 2>/dev/null) )); then
            session_color="${YELLOW}"
        fi
    fi

    local line3_plain="Cost: \$${last_cost} this | \$${session_cost} total | Msgs: ${message_count} | /status"

    if [[ ${#line3_plain} -gt $term_width ]]; then
        local max_len=$((term_width - 3))
        local truncated="${line3_plain:0:$max_len}..."
        line3="${truncated}"
    else
        line3="Cost: ${cost_color}\$${last_cost}${RESET} this | ${session_color}\$${session_cost}${RESET} total | Msgs: ${message_count} | ${DIM}/status${RESET}"
    fi

    # Print all 3 lines using printf for proper color handling
    printf "%b\n" "$line1"
    printf "%b\n" "$line2"
    printf "%b\n" "$line3"
}

# Handle chat message
handle_chat_message() {
    local message="$1"
    local provider model
    provider="$(config_get model_provider)"
    model="$(config_get model_name)"

    # Build message with context using centralized function
    # This ensures chat respects MODE_UPLOADS and context clearing
    local enhanced_message
    enhanced_message=$(build_prompt_with_context "$message" "false")

    # Check if swarm mode should be used
    local response
    local exit_code=0

    if [[ "$SWARM_ENABLED" = "true" ]]; then
        # Estimate tokens to see if swarm is worth it
        local tokens=$(estimate_tokens "$enhanced_message")
        local min_tokens="${SWARM_MIN_TOKENS:-1000}"

        # Force swarm or check if context is large enough
        if [[ "$SWARM_FORCE" = "true" ]] || (( tokens >= min_tokens )); then
            # Show swarm activation
            echo "" >&2
            echo -e "${CYAN}🐝 Activating Swarm Mode${RESET}" >&2
            echo -e "${DIM}Strategy: $SWARM_STRATEGY | Workers: $SWARM_WORKERS | Context: ${tokens} tokens (min: ${min_tokens})${RESET}" >&2
            echo "" >&2

            # Try swarm execution
            if type swarm_execute &>/dev/null; then
                # Swarm execute will output progress to stderr
                response=$(swarm_execute "$enhanced_message" "chat" 2>&1)
                exit_code=$?

                # If swarm failed or returned 1 (fallback), use standard mode
                if [[ $exit_code -ne 0 ]]; then
                    echo -e "${YELLOW}⚠ Swarm mode unavailable, falling back to standard mode${RESET}" >&2
                    echo "" >&2
                    ui_blink "Thinking..."
                    response=$(call_api "$enhanced_message" "$provider" "$model")
                    exit_code=$?
                    ui_clear_line
                fi
            else
                # Swarm not available, use standard
                ui_blink "Thinking..."
                response=$(call_api "$enhanced_message" "$provider" "$model")
                exit_code=$?
                ui_clear_line
            fi
        else
            # Context too small for swarm
            echo -e "${DIM}📝 Using standard mode (context: ${tokens} tokens, need ${min_tokens}+ for swarm)${RESET}" >&2
            ui_blink "Thinking..."
            response=$(call_api "$enhanced_message" "$provider" "$model")
            exit_code=$?
            ui_clear_line
        fi
    else
        # Swarm disabled, use standard mode
        ui_blink "Thinking..."
        response=$(call_api "$enhanced_message" "$provider" "$model")
        exit_code=$?
        ui_clear_line
    fi

    # Check if interrupted (exit code 130)
    if [[ $exit_code -eq 130 ]]; then
        echo ""
        return 130
    fi

    if [[ $exit_code -ne 0 ]]; then
        err "Failed to get response"
        return 1
    fi

    # Display response (redirect to stderr to avoid duplicate output)
    printf "\n" >&2
    printf "%b\n" "${BOLD}AI:${RESET}" >&2
    printf "%b\n" "\033[0;33m${response}${RESET}" >&2  # Simple yellow text
    printf "\n" >&2

    # Track cost first (use enhanced_message to reflect actual tokens sent)
    track_usage "$provider" "$model" "$enhanced_message" "$response"

    # Show status footer (will include the cost just tracked)
    show_status_footer >&2
    echo "" >&2

    # Return only the response text for logging
    echo "$response"
}

# Track API usage and costs
track_usage() {
    local provider="$1"
    local model="$2"
    local input="$3"
    local output="$4"

    local input_tokens output_tokens cost
    input_tokens=$(estimate_tokens "$input")
    output_tokens=$(estimate_tokens "$output")
    cost=$(calculate_cost "$provider" "$model" "$input_tokens" "$output_tokens")

    # Save to usage log
    local workspace logs_dir
    workspace="$(get_workspace)"
    logs_dir="$workspace/logs"
    local usage_log="$logs_dir/usage.jsonl"

    local entry
    entry=$(jq -n \
        --arg provider "$provider" \
        --arg model "$model" \
        --arg in_tokens "$input_tokens" \
        --arg out_tokens "$output_tokens" \
        --arg cost "$cost" \
        --arg timestamp "$(date -Iseconds)" \
        '{provider: $provider, model: $model, input_tokens: ($in_tokens | tonumber), output_tokens: ($out_tokens | tonumber), cost: ($cost | tonumber), timestamp: $timestamp}')

    echo "$entry" >> "$usage_log"

    # Show cost if enabled
    if [[ "$(config_get preferences.show_costs)" == "true" ]]; then
        debug "Cost: \$${cost} USD (${input_tokens} in + ${output_tokens} out tokens)"
    fi
}

# Preflight cost estimation for context operations
preflight_estimate_cost() {
    local operation_type="$1"  # "scan", "load", "context"
    local input_text="${2:-}"
    local provider="${3:-$(config_get model_provider)}"
    local model="${4:-$(config_get model_name)}"

    # Estimate input tokens
    local input_tokens=0
    if [[ -n "$input_text" ]]; then
        # Estimate from provided text
        input_tokens=$(estimate_tokens "$input_text")
    elif type context_state_estimate_tokens &>/dev/null; then
        # Estimate from context files
        input_tokens=$(context_state_estimate_tokens)
    fi

    # If no tokens to estimate, return success
    if [[ $input_tokens -eq 0 ]]; then
        return 0
    fi

    # Estimate output tokens based on operation type
    local output_tokens=0
    case "$operation_type" in
        "scan"|"smartscan")
            # Scan operations: analysis output is usually 0.3x-0.5x input
            output_tokens=$((input_tokens / 2))
            ;;
        "load"|"context")
            # Context load: minimal output expected (just confirmation)
            output_tokens=$((input_tokens / 10))
            ;;
        *)
            # Default: conservative estimate
            output_tokens=$((input_tokens / 2))
            ;;
    esac

    # Calculate estimated cost
    local estimated_cost
    estimated_cost=$(calculate_cost "$provider" "$model" "$input_tokens" "$output_tokens")

    # Get historical costs for comparison (if available)
    local workspace logs_dir usage_log
    workspace="$(get_workspace)"
    logs_dir="$workspace/logs"
    usage_log="$logs_dir/usage.jsonl"

    local historical_avg=""
    if [[ -f "$usage_log" ]] && command -v jq &>/dev/null; then
        # Get average cost from last 5 similar operations
        historical_avg=$(jq -s --arg prov "$provider" \
            'map(select(.provider == $prov)) | .[-5:] | map(.cost) | add / length' \
            "$usage_log" 2>/dev/null || echo "")
    fi

    # Display cost estimate
    echo ""
    echo "${BOLD}${BLUE}💰 Preflight Cost Estimate${RESET}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "  Provider:        $provider"
    echo "  Model:           $model"
    echo "  Operation:       $operation_type"
    echo "  Input tokens:    ~$input_tokens"
    echo "  Output tokens:   ~$output_tokens (estimated)"
    echo "  Estimated cost:  \$${estimated_cost} USD"

    if [[ -n "$historical_avg" && "$historical_avg" != "null" ]]; then
        # Format historical average
        historical_avg=$(awk "BEGIN {printf \"%.6f\", $historical_avg}")
        echo "  Historical avg:  \$${historical_avg} USD (last 5 ops)"

        # Show comparison
        local diff
        diff=$(awk "BEGIN {printf \"%.6f\", $estimated_cost - $historical_avg}")
        if (( $(awk "BEGIN {print ($estimated_cost > $historical_avg)}") )); then
            echo "  Comparison:      ${RED}+\$${diff} higher${RESET}"
        else
            diff=$(awk "BEGIN {printf \"%.6f\", $historical_avg - $estimated_cost}")
            echo "  Comparison:      ${GREEN}-\$${diff} lower${RESET}"
        fi
    fi

    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    # Check if cost estimation is enabled
    if [[ "$(config_get preferences.confirm_before_generate)" != "true" ]]; then
        return 0
    fi

    # Ask for confirmation
    if ! ui_confirm "Proceed with this operation?" "yes"; then
        msg "Operation cancelled"
        return 1
    fi

    return 0
}

# ============================================================================
# TASK COMMANDS
# ============================================================================

# Generate content for task
cmd_generate() {
    local task="${1:-$(config_get current_task)}"
    [[ -z "$task" ]] && task="inbox"

    local workspace tasks_dir task_file
    workspace="$(get_workspace)"
    tasks_dir="$workspace/tasks"
    task_file="$tasks_dir/${task}.prompt.md"

    require_file "$task_file" "Task file"

    msg "Generating content for task: $task"

    # Read base prompt
    local base_prompt
    base_prompt=$(cat "$task_file")

    # Build complete prompt with all context (MODE_UPLOADS, etc.)
    local final_prompt
    final_prompt=$(build_prompt_with_context "$base_prompt" "false")

    # Check if we have context images for vision models
    local context_images=()
    if [[ ${#MODE_UPLOADS[@]} -gt 0 ]]; then
        while IFS= read -r img; do
            [[ -n "$img" ]] && context_images+=("$img")
        done < <(get_context_images)
    fi

    # Estimate first
    if [[ "$(config_get preferences.auto_estimate)" == "true" ]]; then
        estimate_task "$task" "$final_prompt"

        if [[ "$(config_get preferences.confirm_before_generate)" == "true" ]]; then
            if ! ui_confirm "Proceed with generation?"; then
                msg "Cancelled"
                return 0
            fi
        fi
    fi

    # Generate
    local provider model
    provider="$(config_get model_provider)"
    model="$(config_get model_name)"

    echo ""
    msg "Generating with $provider ($model)..."

    # Show context info if present
    if [[ ${#MODE_UPLOADS[@]} -gt 0 ]]; then
        local file_count=${#MODE_UPLOADS[@]}
        local image_count=${#context_images[@]}
        if [[ $image_count -gt 0 ]]; then
            info "Including context: $file_count file(s) ($image_count image(s))"
        else
            info "Including context: $file_count file(s)"
        fi
    fi

    local output
    # Use vision API if we have images
    if [[ ${#context_images[@]} -gt 0 ]]; then
        output=$(call_api_with_images "$final_prompt" "$provider" "$model" "" "${context_images[@]}")
    else
        output=$(call_api "$final_prompt" "$provider" "$model")
    fi

    # Save output
    local output_dir output_file
    output_dir="$workspace/outputs"
    ensure_dir "$output_dir"
    output_file="$output_dir/${task}_$(date +%Y%m%d_%H%M%S).md"

    echo "$output" > "$output_file"

    success "Generated content saved to: $output_file"

    # Show preview
    echo ""
    echo "${CYAN}${BOLD}Preview:${RESET}"
    echo "$output" | head -"$AIWB_CONTEXT_FILE_PREVIEW_LINES"
    if [[ $(echo "$output" | wc -l) -gt $AIWB_CONTEXT_FILE_PREVIEW_LINES ]]; then
        echo "... (truncated, see file for full output)"
    fi

    # Track usage
    track_usage "$provider" "$model" "$prompt" "$output"
}

# Estimate task cost
cmd_estimate() {
    local task="${1:-$(config_get current_task)}"
    [[ -z "$task" ]] && task="inbox"

    local workspace tasks_dir task_file
    workspace="$(get_workspace)"
    tasks_dir="$workspace/tasks"
    task_file="$tasks_dir/${task}.prompt.md"

    require_file "$task_file" "Task file"

    local prompt
    prompt=$(cat "$task_file")

    estimate_task "$task" "$prompt"
}

# Estimate task (helper)
estimate_task() {
    local task="$1"
    local prompt="$2"

    local provider model
    provider="$(config_get model_provider)"
    model="$(config_get model_name)"

    local input_tokens
    input_tokens=$(estimate_tokens "$prompt")

    # Estimate output tokens (rough: 2x input for generation tasks)
    local output_tokens=$((input_tokens * 2))

    # Calculate costs for different tiers
    local tier_mult
    declare -A tiers=(
        ["Basic"]="0.65"
        ["Medium"]="1.00"
        ["Best"]="1.60"
    )

    ui_header "Cost Estimate for: $task"

    echo "Provider: $provider"
    echo "Model:    $model"
    echo ""
    echo "Input tokens:  $input_tokens"
    echo ""

    if $GUM_AVAILABLE; then
        local rows=()
        for tier in Basic Medium Best; do
            mult="${tiers[$tier]}"
            est_out=$(awk "BEGIN {printf \"%.0f\", $output_tokens * $mult}")
            cost=$(calculate_cost "$provider" "$model" "$input_tokens" "$est_out")
            rows+=("${tier}\t${est_out}\t\$${cost}")
        done

        {
            echo "Tier\tOutput Tokens\tCost (USD)"
            printf '%s\n' "${rows[@]}"
        } | column -t -s $'\t'
    else
        printf "%-10s %-15s %s\n" "Tier" "Output Tokens" "Cost (USD)"
        echo "───────────────────────────────────────"
        for tier in Basic Medium Best; do
            mult="${tiers[$tier]}"
            est_out=$(awk "BEGIN {printf \"%.0f\", $output_tokens * $mult}")
            cost=$(calculate_cost "$provider" "$model" "$input_tokens" "$est_out")
            printf "%-10s %-15s \$%s\n" "$tier" "$est_out" "$cost"
        done
    fi
    echo ""
}

# Verify/review generated content
cmd_verify() {
    local task="${1:-$(config_get current_task)}"
    [[ -z "$task" ]] && task="inbox"

    local workspace output_dir
    workspace="$(get_workspace)"
    output_dir="$workspace/outputs"

    # Find latest output for task
    local latest
    latest=$(find "$output_dir" -name "${task}_*.md" -type f 2>/dev/null | sort -r | head -1)

    if [[ -z "$latest" ]]; then
        err "No generated output found for task: $task"
        return 1
    fi

    msg "Verifying: $latest"

    # Read output
    local content prompt
    content=$(cat "$latest")

    # Create verification prompt
    prompt="Review the following generated content and provide constructive feedback:

=== CONTENT TO REVIEW ===
$content

=== INSTRUCTIONS ===
Analyze for:
1. Quality and correctness
2. Completeness
3. Best practices
4. Potential improvements

Provide specific, actionable feedback."

    # Use different provider for verification if available
    local gen_provider verify_provider
    gen_provider="$(config_get model_provider)"

    case "$gen_provider" in
        gemini) verify_provider="claude" ;;
        claude) verify_provider="gemini" ;;
        *) verify_provider="$gen_provider" ;;
    esac

    if has_api_key "$verify_provider"; then
        msg "Using $verify_provider for verification (cross-provider review)"
    else
        verify_provider="$gen_provider"
        msg "Using $gen_provider for verification"
    fi

    # Get model
    local verify_model
    verify_model=$(get_default_model "$verify_provider")

    # Call verifier
    echo ""
    local feedback
    feedback=$(call_api "$prompt" "$verify_provider" "$verify_model")

    # Display feedback
    echo ""
    ui_info_box "VERIFICATION FEEDBACK" "info"
    echo ""
    echo "${CYAN}${BOLD}Feedback:${RESET}"
    echo "$feedback"
    echo ""

    # Save feedback
    local feedback_file="${latest%.md}.feedback.md"
    echo "$feedback" > "$feedback_file"
    success "Feedback saved to: $feedback_file"
}

# Automated Generator-Verifier loop
cmd_refine() {
    local task="${1:-$(config_get current_task)}"
    [[ -z "$task" ]] && task="inbox"

    local iterations="${2:-3}"

    ui_header "Auto-Refine: $task"

    msg "Running Generator-Verifier loop for $iterations iterations..."
    echo ""

    for ((i=1; i<=iterations; i++)); do
        msg "Iteration $i/$iterations"
        echo ""

        # Generate
        ui_style "→ Generating..." "cyan"
        cmd_generate "$task" || break

        # Verify
        echo ""
        ui_style "→ Verifying..." "magenta"
        cmd_verify "$task" || break

        # If not last iteration, update prompt with feedback
        if ((i < iterations)); then
            echo ""
            if ui_confirm "Continue to next iteration?"; then
                # TODO: Incorporate feedback into prompt
                sleep 1
            else
                break
            fi
        fi
    done

    echo ""
    success "Refinement complete!"
}

# ============================================================================
# STATUS AND CONTEXT COMMANDS
# ============================================================================

# ============================================================================
# AUTOMATIC REPO CONTEXT SCANNING
# ============================================================================

# Build repository context automatically

build_repo_context() {
    local mode="${1:-selective}"
    local max_file_size=50000
    local context=""
    local file_count=0

    # Global array to track scanned files (for context persistence)
    declare -g -a SCANNED_FILES=()

    msg "Building $mode context for $(basename "$(pwd)")..."

    # Helper function to check if file is binary by extension
    is_binary_extension() {
        local file="$1"
        case "$file" in
            *.jpg|*.jpeg|*.png|*.gif|*.ico|*.svg|*.webp) return 0 ;;
            *.pdf|*.zip|*.tar|*.gz|*.bz2|*.xz|*.7z) return 0 ;;
            *.exe|*.dll|*.so|*.dylib|*.o|*.a) return 0 ;;
            *.mp4|*.mp3|*.avi|*.mov|*.wav) return 0 ;;
            *.woff|*.woff2|*.ttf|*.eot|*.otf) return 0 ;;
            *.class|*.jar|*.war|*.pyc|*.pyo) return 0 ;;
            *) return 1 ;;
        esac
    }

    if [[ "$mode" == "full" ]]; then
        # Full scan: all text files except in excluded dirs
        msg "Deep scan: reading all source files..."

        # Use process substitution to avoid subshell variable issues
        while IFS= read -r file; do
            # Skip binary files by extension
            if is_binary_extension "$file"; then
                continue
            fi

            # Try to read file content
            local file_content
            file_content=$(head -c "$max_file_size" "$file" 2>/dev/null || true)
            if [[ -n "$file_content" ]]; then
                context+=$'\n'"--- FILE: $file ---"$'\n'
                context+="$file_content"$'\n'
                ((file_count++))
                SCANNED_FILES+=("$(realpath "$file" 2>/dev/null || echo "$file")")
            fi
        done < <(find . -type f \
            -not -path '*/.git/*' \
            -not -path '*/node_modules/*' \
            -not -path '*/dist/*' \
            -not -path '*/build/*' \
            -not -path '*/__pycache__/*' \
            -not -path '*/.next/*' \
            -not -path '*/vendor/*' \
            -not -path '*/target/*' \
            -not -path '*/.venv/*' \
            -not -path '*/venv/*' \
            2>/dev/null)

    else
        # Selective: key files only
        msg "Smart scan: reading key files..."

        # Config files
        for pattern in "package.json" "tsconfig.json" "vite.config.*" "*.config.js" "*.config.ts" "Cargo.toml" "go.mod" "requirements.txt" ".env.example"; do
            while IFS= read -r file; do
                local file_content
                file_content=$(head -c "$max_file_size" "$file" 2>/dev/null || true)
                if [[ -n "$file_content" ]]; then
                    context+=$'\n'"--- FILE: $file ---"$'\n'
                    context+="$file_content"$'\n'
                    ((file_count++))
                    SCANNED_FILES+=("$(realpath "$file" 2>/dev/null || echo "$file")")
                fi
            done < <(find . -maxdepth 2 -type f -name "$pattern" \
                -not -path '*/.git/*' \
                -not -path '*/node_modules/*' \
                2>/dev/null)
        done

        # Documentation
        for pattern in "README*" "CHANGELOG*" "LICENSE*"; do
            while IFS= read -r file; do
                local file_content
                file_content=$(head -c "$max_file_size" "$file" 2>/dev/null || true)
                if [[ -n "$file_content" ]]; then
                    context+=$'\n'"--- FILE: $file ---"$'\n'
                    context+="$file_content"$'\n'
                    ((file_count++))
                    SCANNED_FILES+=("$(realpath "$file" 2>/dev/null || echo "$file")")
                fi
            done < <(find . -maxdepth 2 -type f -name "$pattern" \
                -not -path '*/.git/*' \
                2>/dev/null)
        done

        # Main source files - NO MORE file command!
        for pattern in "main.*" "index.*" "app.*" "App.*" "server.*" "__init__.py" "mod.rs"; do
            while IFS= read -r file; do
                # Skip binary extensions
                if is_binary_extension "$file"; then
                    continue
                fi

                local file_content
                file_content=$(head -c "$max_file_size" "$file" 2>/dev/null || true)
                if [[ -n "$file_content" ]]; then
                    SCANNED_FILES+=("$(realpath "$file" 2>/dev/null || echo "$file")")
                    context+=$'\n'"--- FILE: $file ---"$'\n'
                    context+="$file_content"$'\n'
                    ((file_count++))
                fi
            done < <(find . -maxdepth 3 -type f -name "$pattern" \
                -not -path '*/.git/*' \
                -not -path '*/node_modules/*' \
                -not -path '*/dist/*' \
                2>/dev/null)
        done
    fi

    msg "Context built: $file_count files scanned"
    echo "$context"
}

# Scan repository - full deep scan
cmd_scanrepo() {
    ui_header "Full Repository Scan"

    # Always enable repo mode for current directory
    detect_repo

    echo "Folder: ${AIWB_REPO_NAME}"
    echo "Path: ${AIWB_REPO_PATH}"
    if [[ -n "$AIWB_REPO_REMOTE" ]]; then
        echo "Git: ${AIWB_REPO_BRANCH} (${AIWB_REPO_REMOTE})"
    else
        echo "Type: Local directory (not a git repository)"
    fi
    echo ""

    # Run security audit for git repositories
    if [[ -n "$AIWB_REPO_REMOTE" ]]; then
        if type validate_gitignore &>/dev/null; then
            validate_gitignore
            echo ""
        fi
        if type audit_git_exposure &>/dev/null; then
            audit_git_exposure
            echo ""
        fi
    fi

    echo "This will scan ALL source files in this folder."
    echo ""

    if ! ui_confirm "Proceed with full scan?" "yes"; then
        msg "Cancelled"
        return 0
    fi

    # Build full context
    local context
    context=$(build_repo_context "full")

    # Preflight cost estimation
    if ! preflight_estimate_cost "scan" "$context"; then
        return 0
    fi
    
    if [[ -z "$context" ]]; then
        warn "No files found to scan"
        return 1
    fi
    
    # Create analysis prompt
    local prompt="Analyze this repository and provide a comprehensive overview:

Repository: ${AIWB_REPO_NAME}
Branch: ${AIWB_REPO_BRANCH}

$context

Please provide:
1. Project type and purpose
2. Technology stack and dependencies
3. Architecture and key components
4. Main features and functionality
5. Code quality observations
6. Suggestions for improvements"
    
    # Get provider and model
    local provider model
    provider="$(config_get model_provider)"
    model="$(config_get model_name)"
    
    echo ""
    msg "Analyzing with $provider ($model)..."
    echo ""
    
    # Call API
    local analysis
    analysis=$(call_api "$prompt" "$provider" "$model")
    
    # Display results
    echo ""
    ui_header "Repository Analysis"
    echo "$analysis"
    echo ""
    
    # Save analysis
    local workspace output_dir
    workspace="$(get_workspace)"
    output_dir="$workspace/outputs"
    ensure_dir "$output_dir"
    
    local analysis_file="$output_dir/repo_analysis_$(date +%Y%m%d_%H%M%S).md"
    echo "$analysis" > "$analysis_file"

    success "Analysis saved to: $analysis_file"

    # Track usage
    track_usage "$provider" "$model" "$prompt" "$analysis"

    # Save to context state (if context_state functions are available)
    if type context_state_save_scan &>/dev/null; then
        # Count files from context (count "--- FILE:" markers)
        local scan_file_count
        scan_file_count=$(echo "$context" | grep -c "^--- FILE:" || echo "0")
        context_state_save_scan "full" "$analysis_file" "$scan_file_count"

        # Also save all scanned source files to context for persistence
        if type context_state_add_file &>/dev/null && [[ ${#SCANNED_FILES[@]} -gt 0 ]]; then
            for scanned_file in "${SCANNED_FILES[@]}"; do
                context_state_add_file "$scanned_file" "scan" 2>/dev/null || true
            done

            # Auto-load context into current session (no need for manual /contextload)
            if type context_state_load_into_mode &>/dev/null; then
                echo ""
                context_state_load_into_mode
            else
                msg "Context saved: ${#SCANNED_FILES[@]} source files + analysis (run /contextload to reload)"
            fi
        else
            msg "Context saved for future use (run /contextshow to view)"
        fi
    fi
}

# Smart scan - selective key files only
cmd_smartscan() {
    ui_header "Smart Repository Scan"

    # Always enable repo mode for current directory (works for ANY folder)
    detect_repo

    echo "Folder: ${AIWB_REPO_NAME}"
    echo "Path: ${AIWB_REPO_PATH}"
    if [[ -n "$AIWB_REPO_REMOTE" ]]; then
        echo "Git: ${AIWB_REPO_BRANCH} (${AIWB_REPO_REMOTE})"
    else
        echo "Type: Local directory (not a git repository)"
    fi
    echo ""

    # Run security audit for git repositories
    if [[ -n "$AIWB_REPO_REMOTE" ]]; then
        if type validate_gitignore &>/dev/null; then
            validate_gitignore
            echo ""
        fi
        if type audit_git_exposure &>/dev/null; then
            audit_git_exposure
            echo ""
        fi
    fi

    echo "This will scan configuration files, documentation, and main entry points."
    echo ""

    # Build selective context
    local context
    context=$(build_repo_context "selective")

    if [[ -z "$context" ]]; then
        warn "No key files found to scan"
        return 1
    fi

    # Preflight cost estimation
    if ! preflight_estimate_cost "smartscan" "$context"; then
        return 0
    fi

    # Create analysis prompt
    local prompt="Provide a quick overview of this folder based on key files:

Folder: ${AIWB_REPO_NAME}
${AIWB_REPO_BRANCH:+Branch: $AIWB_REPO_BRANCH}

$context

Please summarize:
1. Project type and main purpose
2. Key technologies used
3. Main dependencies
4. Entry points and structure"

    # Get provider and model
    local provider model
    provider="$(config_get model_provider)"
    model="$(config_get model_name)"

    echo ""
    msg "Analyzing with $provider ($model)..."
    echo ""

    # Call API
    local analysis
    analysis=$(call_api "$prompt" "$provider" "$model")

    # Display results
    echo ""
    ui_header "Quick Analysis"
    echo "$analysis"
    echo ""

    # Save analysis
    local workspace output_dir
    workspace="$(get_workspace)"
    output_dir="$workspace/outputs"
    ensure_dir "$output_dir"

    local analysis_file="$output_dir/smart_scan_$(date +%Y%m%d_%H%M%S).md"
    echo "$analysis" > "$analysis_file"

    success "Analysis saved to: $analysis_file"

    # Track usage
    track_usage "$provider" "$model" "$prompt" "$analysis"

    # Save to context state (if context_state functions are available)
    if type context_state_save_scan &>/dev/null; then
        # Count files from context (count "--- FILE:" markers)
        local scan_file_count
        scan_file_count=$(echo "$context" | grep -c "^--- FILE:" || echo "0")
        context_state_save_scan "selective" "$analysis_file" "$scan_file_count"

        # Also save all scanned source files to context for persistence
        if type context_state_add_file &>/dev/null && [[ ${#SCANNED_FILES[@]} -gt 0 ]]; then
            for scanned_file in "${SCANNED_FILES[@]}"; do
                context_state_add_file "$scanned_file" "scan" 2>/dev/null || true
            done

            # Auto-load context into current session (no need for manual /contextload)
            if type context_state_load_into_mode &>/dev/null; then
                echo ""
                context_state_load_into_mode
            else
                msg "Context saved: ${#SCANNED_FILES[@]} source files + analysis (run /contextload to reload)"
            fi
        else
            msg "Context saved for future use (run /contextshow to view)"
        fi
    fi
}

# ============================================================================
# CONTEXT MANAGEMENT COMMANDS
# ============================================================================

# Load saved context into current session
cmd_contextload() {
    local workspace
    workspace=$(get_workspace)

    if ! context_state_exists; then
        warn "No saved context found"
        echo ""
        info "Tip: Run /scanrepo or /smartscan to build context first"
        return 1
    fi

    # Check age and warn if stale
    context_state_check_age

    # Show what will be loaded
    echo ""
    context_state_show
    echo ""

    # Preflight cost estimation
    if ! preflight_estimate_cost "load"; then
        return 0
    fi

    # Load context files into MODE_UPLOADS
    context_state_load_into_mode

    # Show cost footer after loading
    local tokens
    tokens=$(context_state_estimate_tokens)
    if [[ $tokens -gt 0 ]]; then
        local provider model estimated_cost
        provider="$(config_get model_provider)"
        model="$(config_get model_name)"
        # Estimate minimal cost for having context loaded (output = 10% of input)
        estimated_cost=$(calculate_cost "$provider" "$model" "$tokens" "$((tokens / 10))")

        echo ""
        echo "${DIM}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
        echo "${DIM}  Context loaded: ~${tokens} tokens${RESET}"
        echo "${DIM}  Estimated cost per use: ~\$${estimated_cost} USD${RESET}"
        echo "${DIM}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${RESET}"
    fi

    echo ""
    success "Context loaded. Use /make, /debug, or /tweak to use it."
}

# Save current context (MODE_UPLOADS) to state file
cmd_contextsave() {
    if [[ ${#MODE_UPLOADS[@]} -eq 0 ]]; then
        warn "No files in context to save"
        echo ""
        info "Tip: Use /context to add files, or run /scanrepo first"
        return 1
    fi

    # Save MODE_UPLOADS to context state
    context_state_save_from_mode

    # Show summary
    echo ""
    context_state_show
}

# Display current context state
cmd_contextshow() {
    if ! context_state_exists; then
        warn "No active context state"
        echo ""
        info "Tip: Run /scanrepo or add files with /context to create context"
        return 1
    fi

    context_state_show

    # Also check size
    echo ""
    context_state_check_size
}

# Clear all context and start fresh
cmd_contextclear() {
    if ! context_state_exists && [[ ${#MODE_UPLOADS[@]} -eq 0 ]] && [[ -z "$MODE_PROMPT" ]] && [[ -z "$MODE_INSTRUCT_FILE" ]]; then
        warn "No context to clear"
        return 0
    fi

    # Show what will be cleared
    warn "This will clear:"
    echo "   • All context files (persistent and in-memory)"
    echo "   • Conversation history"
    echo "   • Scan results"
    echo "   • Mode prompts and instructions"
    echo "   • Check/verify settings"

    # Show current context
    local total_items=0
    if [[ ${#MODE_UPLOADS[@]} -gt 0 ]]; then
        echo ""
        echo "   In-memory context ($([[ ${#MODE_UPLOADS[@]} -gt 0 ]] && echo ${#MODE_UPLOADS[@]} || echo 0) file(s)):"
        for item in "${MODE_UPLOADS[@]}"; do
            echo "     - $(basename "$item")"
            ((total_items++))
        done
    fi

    if context_state_exists; then
        local persistent_count=$(context_state_get "context_files | length" 2>/dev/null || echo "0")
        if [[ "$persistent_count" != "0" && "$persistent_count" != "null" ]]; then
            echo ""
            echo "   Persistent context ($persistent_count file(s))"
            ((total_items+=persistent_count))
        fi
    fi

    echo ""
    echo "   Total items to clear: $total_items"
    echo ""

    if ! ui_confirm "Clear all context?" "no"; then
        msg "Cancelled"
        return 0
    fi

    # Clear persistent context state
    if context_state_exists; then
        context_state_clear
    fi

    # Clear all MODE state variables using centralized function
    clear_all_mode_state

    success "All context cleared successfully"
    info "You can now start fresh or load new context"
}

# Re-scan repository and update context
cmd_contextrefresh() {
    echo "🔄 Refreshing context..."
    echo ""

    # Detect current scan type from context state
    local scan_type="full"

    if context_state_exists; then
        scan_type=$(context_state_get "last_scan.type")
        if [[ -z "$scan_type" || "$scan_type" == "null" ]]; then
            scan_type="full"
        fi
    fi

    echo "Previous scan type: $scan_type"
    echo ""

    # Prompt for scan type
    if $GUM_AVAILABLE; then
        scan_type=$(gum choose "full" "selective" --header "Choose scan type:")
    else
        echo "Choose scan type:"
        echo "  1) full     - Scan entire repository"
        echo "  2) selective - Scan key files only"
        read -p "Enter choice [1-2]: " -r choice
        case "$choice" in
            1) scan_type="full" ;;
            2) scan_type="selective" ;;
            *) scan_type="full" ;;
        esac
    fi

    # Clear old context first (completely)
    if context_state_exists; then
        context_state_clear
    fi
    clear_all_mode_state

    # Run appropriate scan
    if [[ "$scan_type" == "selective" ]]; then
        cmd_smartscan
    else
        cmd_scanrepo
    fi

    echo ""
    success "Context refreshed"
    context_state_show
}

# Remove specific file from context
cmd_contextremove() {
    if ! context_state_exists; then
        warn "No context to remove from"
        return 1
    fi

    # Get list of context files
    local files
    mapfile -t files < <(context_state_list_files)

    if [[ ${#files[@]} -eq 0 ]]; then
        warn "No files in context"
        return 1
    fi

    # Interactive selection
    local file_to_remove

    if $GUM_AVAILABLE; then
        file_to_remove=$(printf '%s\n' "${files[@]}" | gum filter --placeholder "Select file to remove")
    else
        echo "Files in context:"
        local i=1
        for file in "${files[@]}"; do
            echo "  $i) $file"
            ((i++))
        done
        echo ""
        read -p "Enter number to remove: " -r choice
        file_to_remove="${files[$((choice-1))]}"
    fi

    if [[ -z "$file_to_remove" ]]; then
        msg "Cancelled"
        return 0
    fi

    # Remove from context state
    context_state_remove_file "$file_to_remove"

    # Also remove from MODE_UPLOADS if present
    local new_uploads=()
    for file in "${MODE_UPLOADS[@]}"; do
        if [[ "$file" != "$file_to_remove" ]]; then
            new_uploads+=("$file")
        fi
    done
    MODE_UPLOADS=("${new_uploads[@]}")

    success "Removed: $file_to_remove"
}

# Swarm mode command
cmd_swarm() {
    local action="${1:-menu}"

    case "$action" in
        menu|"")
            # Interactive menu (default)
            menu_swarm
            ;;
        status)
            # Show detailed swarm status
            ui_header "Swarm Mode Status"
            echo "Enabled: $SWARM_ENABLED"
            echo "Strategy: $SWARM_STRATEGY"
            echo "Workers: $SWARM_WORKERS"
            echo "Worker Provider: $SWARM_WORKER_PROVIDER"
            echo "Worker Model: $SWARM_WORKER_MODEL"
            echo "Aggregator Provider: $SWARM_AGGREGATOR_PROVIDER"
            echo "Aggregator Model: $SWARM_AGGREGATOR_MODEL"
            echo ""
            echo "Display: $(get_swarm_display)"
            ;;
        on|enable)
            # Enable swarm mode
            if [[ "$SWARM_ENABLED" = "true" ]]; then
                info "Swarm mode is already enabled"
            else
                export SWARM_ENABLED="true"
                config_set "swarm.enabled" "true"
                success "Swarm mode ENABLED 🐝"
            fi
            ;;
        off|disable)
            # Disable swarm mode
            if [[ "$SWARM_ENABLED" = "false" ]]; then
                info "Swarm mode is already disabled"
            else
                export SWARM_ENABLED="false"
                config_set "swarm.enabled" "false"
                msg "Swarm mode DISABLED"
            fi
            ;;
        toggle)
            # Toggle swarm mode
            swarm_toggle
            ;;
        *)
            warn "Unknown swarm action: $action"
            echo "Usage: /swarm [menu|status|on|off|toggle]"
            ;;
    esac
}

# Overall status overview
cmd_status_overview() {
    ui_header "AIWB Status Overview"

    # Current configuration
    echo "Configuration:"
    echo "  Provider: $(config_get model_provider)"
    echo "  Model: $(config_get model_name)"
    echo "  Workspace: $(get_workspace)"
    echo ""

    # Mode states
    echo "Mode States:"
    echo "  Make - Model: ${MODE_MODEL_PROVIDER:-default} ${MODE_MODEL_NAME:-default}"
    echo "  Tweak - Model: ${MODE_MODEL_PROVIDER:-default} ${MODE_MODEL_NAME:-default}"
    echo "  Debug - Model: ${MODE_MODEL_PROVIDER:-default} ${MODE_MODEL_NAME:-default}"
    echo ""

    # Context
    echo "Global Context:"

    # Show persistent context (same as footer)
    local context_summary=$(get_context_summary)
    echo "  Persistent: $context_summary"

    # Show in-memory context
    if [[ ${#MODE_UPLOADS[@]} -gt 0 ]]; then
        echo "  In-Memory:"
        for item in "${MODE_UPLOADS[@]}"; do
            echo "    - $item"
        done
    else
        echo "  In-Memory: (no files loaded)"
    fi
    echo ""

    # Costs
    local workspace logs_dir usage_log
    workspace="$(get_workspace)"
    logs_dir="$workspace/logs"
    usage_log="$logs_dir/usage.jsonl"

    if [[ -f "$usage_log" ]]; then
        echo "Estimated Costs (All Time):"
        echo "${DIM}  (Note: Estimates may differ from actual API costs)${RESET}"
        local total
        total=$(jq -s 'map(.cost) | add' "$usage_log" 2>/dev/null || echo "0")
        total=$(printf "%.2f" "$total")
        echo "  Total estimated: \$${total} USD"

        echo ""
        echo "  By Provider:"
        jq -s 'group_by(.provider) | map({provider: .[0].provider, cost: (map(.cost) | add)}) | .[]' "$usage_log" 2>/dev/null | \
            jq -r '"    \(.provider): $\(.cost | tonumber | . * 100 | round / 100)"'
    else
        echo "Costs:"
        echo "  No usage data yet"
    fi
    echo ""

    # Recent activity
    echo "Recent Activity:"
    if [[ -d "$workspace/outputs" ]]; then
        local recent
        recent=$(find "$workspace/outputs" -name "*.md" ! -name "*.feedback.md" -type f 2>/dev/null | sort -r | head -3)
        if [[ -n "$recent" ]]; then
            echo "$recent" | while read -r file; do
                local basename=$(basename "$file")
                local timestamp=$(stat -f%Sm -t "%Y-%m-%d %H:%M" "$file" 2>/dev/null || stat -c%y "$file" 2>/dev/null | cut -d'.' -f1)
                echo "  - $basename ($timestamp)"
            done
        else
            echo "  (no recent outputs)"
        fi
    else
        echo "  (no outputs yet)"
    fi
    echo ""
}

# Context management - interactive menu
cmd_context() {
    while true; do
        local context_count=${#MODE_UPLOADS[@]}

        # Also check persistent context
        local persistent_count=0
        local has_scan=false
        if type context_state_exists &>/dev/null && context_state_exists; then
            # Count manual AND scan files (all types)
            persistent_count=$(context_state_get "context_files" 2>/dev/null | jq 'length' 2>/dev/null || echo 0)
            local scan_file
            scan_file=$(context_state_get "last_scan.output_file" 2>/dev/null)
            [[ -n "$scan_file" && -f "$scan_file" ]] && has_scan=true
        fi

        local total_context=$((context_count + persistent_count))
        local context_display="$context_count in-memory"
        if [[ $persistent_count -gt 0 ]] || [[ "$has_scan" == "true" ]]; then
            context_display="$context_count in-memory, $persistent_count saved"
            if [[ "$has_scan" == "true" ]]; then
                context_display="$context_display + scan"
            fi
        fi

        local choice
        choice=$(ui_choose "Context Manager ($context_display)" \
            "Add files/directories" \
            "List current context (in-memory)" \
            "Show persistent context (saved state)" \
            "Load persistent context into memory" \
            "Save in-memory context to persistent state" \
            "Remove from in-memory context" \
            "Clear all context" \
            "Back")

        case "$choice" in
            "Add files/directories")
                menu_uploads
                ;;
            "List current context (in-memory)")
                echo ""
                echo "In-Memory Context (MODE_UPLOADS):"
                echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
                if [[ ${#MODE_UPLOADS[@]} -gt 0 ]]; then
                    for item in "${MODE_UPLOADS[@]}"; do
                        if [[ -f "$item" ]]; then
                            local size=$(stat -f%z "$item" 2>/dev/null || stat -c%s "$item" 2>/dev/null || echo 0)
                            local size_kb=$((size / 1024))
                            if is_image_file "$item"; then
                                echo "  - [IMAGE] $item (${size_kb} KB)"
                            else
                                echo "  - [FILE] $item (${size_kb} KB)"
                            fi
                        elif [[ -d "$item" ]]; then
                            local count=$(find "$item" -type f 2>/dev/null | wc -l | tr -d ' ')
                            echo "  - [DIR] $item/ ($count files)"
                        fi
                    done
                    echo ""
                    echo "Total: ${#MODE_UPLOADS[@]} item(s) in memory"
                else
                    echo "  (no in-memory context files)"
                    echo ""
                    echo "💡 To add context: /context → Add files/directories"
                    echo "   Or load from saved state: /context → Load persistent context"
                fi
                echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
                ui_pause
                ;;
            "Show persistent context (saved state)")
                echo ""
                if type context_state_show &>/dev/null; then
                    context_state_show
                else
                    warn "Context persistence not available"
                fi
                echo ""
                ui_pause
                ;;
            "Load persistent context into memory")
                if type context_state_load_into_mode &>/dev/null; then
                    echo ""
                    context_state_load_into_mode
                    echo ""
                    ui_pause
                else
                    warn "Context persistence not available"
                fi
                ;;
            "Save in-memory context to persistent state")
                if type context_state_save_from_mode &>/dev/null; then
                    echo ""
                    echo "Saving in-memory context to persistent state..."
                    context_state_save_from_mode
                    echo ""
                    ui_pause
                else
                    warn "Context persistence not available"
                fi
                ;;
            "Remove from in-memory context")
                if [[ ${#MODE_UPLOADS[@]} -eq 0 ]]; then
                    warn "No in-memory context files to remove"
                else
                    # Build list of items for removal
                    local items=()
                    for item in "${MODE_UPLOADS[@]}"; do
                        items+=("$item")
                    done
                    items+=("Back")

                    local to_remove
                    to_remove=$(ui_choose "Select file to remove" "${items[@]}")

                    if [[ -n "$to_remove" ]] && [[ "$to_remove" != "Back" ]]; then
                        local new_uploads=()
                        for upload in "${MODE_UPLOADS[@]}"; do
                            if [[ "$upload" != "$to_remove" ]]; then
                                new_uploads+=("$upload")
                            fi
                        done
                        MODE_UPLOADS=("${new_uploads[@]}")
                        success "Removed from in-memory context: $to_remove"
                    fi
                fi
                ;;
            "Clear all context")
                if [[ ${#MODE_UPLOADS[@]} -gt 0 ]] || [[ $persistent_count -gt 0 ]]; then
                    local choices=()
                    [[ ${#MODE_UPLOADS[@]} -gt 0 ]] && choices+=("Clear in-memory context only")
                    if type context_state_exists &>/dev/null && context_state_exists; then
                        choices+=("Clear persistent context only")
                        choices+=("Clear ALL (recommended)")
                    fi
                    choices+=("Cancel")

                    local clear_choice
                    clear_choice=$(ui_choose "What to clear?" "${choices[@]}")

                    case "$clear_choice" in
                        "Clear in-memory context only")
                            clear_all_mode_state
                            success "In-memory context and mode state cleared"
                            warn "Note: Persistent context still exists. Use 'Clear ALL' for complete clearing."
                            ;;
                        "Clear persistent context only")
                            if type context_state_clear &>/dev/null; then
                                context_state_clear
                            fi
                            warn "Note: In-memory context still loaded. Use 'Clear ALL' for complete clearing."
                            ;;
                        "Clear ALL (recommended)")
                            # Clear persistent state
                            if type context_state_clear &>/dev/null && context_state_exists; then
                                context_state_clear
                            fi
                            # Clear all in-memory state
                            clear_all_mode_state
                            success "All context cleared successfully (persistent + in-memory)"
                            info "You can now start fresh or load new context"
                            ;;
                    esac
                else
                    info "Context is already empty"
                fi
                ;;
            "Back"|"")
                break
                ;;
        esac
    done
}

# ============================================================================
# QUICK WORKFLOW COMMANDS
# ============================================================================

# Quick one-shot generation with auto-verify
cmd_quick() {
    local description="$*"

    # If no description provided, prompt for it
    if [[ -z "$description" ]]; then
        description=$(ui_input "What do you want to create?")
        [[ -z "$description" ]] && return 0
    fi

    ui_header "Quick Generation"
    msg "Creating: $description"
    echo ""

    # Create temporary task
    local task_name="quick_$(date +%s)"
    local workspace tasks_dir
    workspace="$(get_workspace)"
    tasks_dir="$workspace/tasks"
    ensure_dir "$tasks_dir"

    local task_file="$tasks_dir/${task_name}.prompt.md"
    echo "$description" > "$task_file"

    # Set as current task
    config_set current_task "$task_name"

    # Check if should auto-verify
    local auto_verify="true"
    if [[ "$(config_get preferences.confirm_before_generate)" == "true" ]]; then
        if ! ui_confirm "Auto-verify with a different model?" "yes"; then
            auto_verify="false"
        fi
    fi

    # Generate
    cmd_generate "$task_name"

    # Auto-verify if enabled
    if [[ "$auto_verify" == "true" ]]; then
        echo ""
        msg "Auto-verifying with cross-provider review..."
        cmd_verify "$task_name"
    fi

    # Cleanup temporary task
    rm -f "$task_file"

    echo ""
    success "Done! Check outputs with: ./aiwb history"
}

# Interactive wizard workflow
cmd_wizard() {
    clear 2>/dev/null || true
    ui_header "AIWB Wizard"

    # Step 1: What to build
    local project_type
    project_type=$(ui_choose "What do you want to build?" \
        "Simple script/tool" \
        "Web application" \
        "API integration" \
        "Data processing pipeline" \
        "Documentation" \
        "Just chat/brainstorm" \
        "Cancel")

    [[ "$project_type" == "Cancel" || -z "$project_type" ]] && return 0

    # Step 2: Get description
    echo ""
    local description
    description=$(ui_input "Brief description of your project:")
    [[ -z "$description" ]] && return 0

    # Step 3: Choose language/technology (if applicable)
    local tech=""
    if [[ "$project_type" != "Just chat/brainstorm" && "$project_type" != "Documentation" ]]; then
        echo ""
        tech=$(ui_choose "Choose technology/language:" \
            "Bash/Shell" \
            "Python" \
            "JavaScript/Node.js" \
            "Go" \
            "Rust" \
            "Other/Mixed" \
            "Skip")
        [[ "$tech" == "Skip" ]] && tech=""
    fi

    # Step 4: Special requirements
    echo ""
    local requirements
    requirements=$(ui_input "Any special requirements? (optional)" "" "e.g., security, performance, compatibility")

    # Step 5: Build the prompt
    local full_prompt="$description"
    [[ -n "$tech" && "$tech" != "Skip" ]] && full_prompt="$full_prompt

Technology: $tech"
    [[ -n "$requirements" ]] && full_prompt="$full_prompt

Requirements:
$requirements"

    # Step 6: Show preview and confirm
    echo ""
    ui_header "Preview"
    echo "$full_prompt"
    echo ""

    if ! ui_confirm "Proceed with generation?" "yes"; then
        msg "Cancelled"
        return 0
    fi

    # Step 7: Choose workflow
    echo ""
    local workflow
    workflow=$(ui_choose "Choose workflow:" \
        "Quick (generate + verify)" \
        "Generate only" \
        "Full refinement loop (3 iterations)")

    # Execute
    echo ""
    case "$workflow" in
        "Quick"*)
            cmd_quick "$full_prompt"
            ;;
        "Generate only")
            local task_name="wizard_$(date +%s)"
            local workspace tasks_dir
            workspace="$(get_workspace)"
            tasks_dir="$workspace/tasks"
            ensure_dir "$tasks_dir"
            echo "$full_prompt" > "$tasks_dir/${task_name}.prompt.md"
            config_set current_task "$task_name"
            cmd_generate "$task_name"
            rm -f "$tasks_dir/${task_name}.prompt.md"
            ;;
        "Full refinement"*)
            local task_name="wizard_$(date +%s)"
            local workspace tasks_dir
            workspace="$(get_workspace)"
            tasks_dir="$workspace/tasks"
            ensure_dir "$tasks_dir"
            echo "$full_prompt" > "$tasks_dir/${task_name}.prompt.md"
            config_set current_task "$task_name"
            cmd_refine "$task_name" 3
            rm -f "$tasks_dir/${task_name}.prompt.md"
            ;;
    esac
}

# Make (alias for quick)
cmd_make() {
    cmd_quick "$@"
}

# Improve last output
cmd_improve() {
    local suggestion="$*"

    if [[ -z "$suggestion" ]]; then
        suggestion=$(ui_input "What improvements do you want?")
        [[ -z "$suggestion" ]] && return 0
    fi

    # Find last output
    local workspace output_dir
    workspace="$(get_workspace)"
    output_dir="$workspace/outputs"

    local latest
    latest=$(find "$output_dir" -name "*.md" -type f ! -name "*.feedback.md" 2>/dev/null | sort -r | head -1)

    if [[ -z "$latest" ]]; then
        err "No previous output found to improve"
        return 1
    fi

    msg "Improving: $(basename "$latest")"
    echo ""

    # Read last output
    local content
    content=$(cat "$latest")

    # Create improvement prompt
    local prompt="Improve the following based on this feedback: $suggestion

=== ORIGINAL CONTENT ===
$content

=== IMPROVEMENT REQUEST ===
$suggestion

Please provide an improved version incorporating the feedback."

    # Generate improvement
    local provider model
    provider="$(config_get model_provider)"
    model="$(config_get model_name)"

    msg "Generating improvements with $provider..."
    echo ""

    local improved
    improved=$(call_api "$prompt" "$provider" "$model")

    # Save improved version
    local improved_file="${latest%.md}_improved_$(date +%s).md"
    echo "$improved" > "$improved_file"

    success "Improved version saved to: $improved_file"

    # Show preview
    echo ""
    echo "${CYAN}${BOLD}Preview:${RESET}"
    echo "$improved" | head -20

    # Track usage
    track_usage "$provider" "$model" "$prompt" "$improved"
}

# Browse and use templates
cmd_templates() {
    local workspace templates_dir
    workspace="$(get_workspace)"
    templates_dir="$workspace/templates"

    # Create templates directory if it doesn't exist
    ensure_dir "$templates_dir"

    # Create default templates if none exist
    if [[ -z "$(ls -A "$templates_dir" 2>/dev/null)" ]]; then
        create_default_templates "$templates_dir"
    fi

    # List templates
    local templates
    templates=($(find "$templates_dir" -name "*.md" -type f -exec basename {} .md \; 2>/dev/null | sort))

    if [[ ${#templates[@]} -eq 0 ]]; then
        info "No templates found. Creating defaults..."
        create_default_templates "$templates_dir"
        templates=($(find "$templates_dir" -name "*.md" -type f -exec basename {} .md \; 2>/dev/null | sort))
    fi

    # Choose template
    local template
    template=$(ui_choose "Select Template" "${templates[@]}" "Cancel")

    [[ "$template" == "Cancel" || -z "$template" ]] && return 0

    # Load template
    local template_file="$templates_dir/${template}.md"
    local template_content
    template_content=$(cat "$template_file")

    # Show template
    echo ""
    ui_header "Template: $template"
    echo "$template_content"
    echo ""

    if ui_confirm "Use this template?" "yes"; then
        cmd_quick "$template_content"
    fi
}

# Create default templates
create_default_templates() {
    local templates_dir="$1"

    cat > "$templates_dir/bash-script.md" <<'EOF'
Create a bash script with the following requirements:
- POSIX-compliant for maximum compatibility
- Proper error handling with set -euo pipefail
- Clear usage/help message
- Input validation
- Colorful output for better UX
- Well-commented code

Describe what the script should do:
EOF

    cat > "$templates_dir/python-cli.md" <<'EOF'
Create a Python CLI application with:
- argparse for argument parsing
- Proper error handling
- Type hints
- Docstrings
- Logging support
- Clean project structure

Describe the CLI tool's purpose:
EOF

    cat > "$templates_dir/api-integration.md" <<'EOF'
Create an API integration with:
- Proper authentication handling
- Error handling and retries
- Rate limiting consideration
- Request/response validation
- Clear documentation

Describe the API to integrate:
EOF

    cat > "$templates_dir/data-pipeline.md" <<'EOF'
Create a data processing pipeline with:
- Input validation
- Error handling
- Progress reporting
- Data transformation steps
- Output formatting

Describe the data processing requirements:
EOF

    success "Created default templates in $templates_dir"
}

# ============================================================================
# CONFIGURATION COMMANDS
# ============================================================================

# Manage API keys
cmd_keys() {
    setup_keys_interactive
}

# Models menu
cmd_models() {
    while true; do
        local provider model
        provider="$(config_get model_provider)"
        model="$(config_get model_name)"

        local choice
        choice=$(ui_choose "Settings Menu" \
            "Provider: $provider" \
            "Model: $model" \
            "Preferences" \
            "Cost Tracking" \
            "Security" \
            "Back")

        case "$choice" in
            "Provider:"*)
                local new_provider
                new_provider=$(ui_choose "Select Provider" "gemini" "claude" "openai" "groq" "xai" "ollama")
                [[ -n "$new_provider" ]] && {
                    config_set model_provider "$new_provider"
                    config_set model_name "$(get_default_model "$new_provider")"
                    success "Provider set to: $new_provider"
                }
                ;;
            "Model:"*)
                provider="$(config_get model_provider)"
                local models available
                available=$(get_available_models "$provider")
                local new_model
                new_model=$(ui_choose "Select Model" $available)
                [[ -n "$new_model" ]] && {
                    config_set model_name "$new_model"
                    success "Model set to: $new_model"
                }
                ;;
            "Preferences")
                settings_preferences
                ;;
            "Cost Tracking")
                settings_cost_tracking
                ;;
            "Security")
                settings_security
                ;;
            "Back"|"")
                break
                ;;
        esac
    done

    save_session
}

# Preferences settings
settings_preferences() {
    while true; do
        local auto_estimate confirm_gen show_costs
        auto_estimate="$(config_get preferences.auto_estimate)"
        confirm_gen="$(config_get preferences.confirm_before_generate)"
        show_costs="$(config_get preferences.show_costs)"

        local choice
        choice=$(ui_choose "Preferences" \
            "Auto-estimate: $auto_estimate" \
            "Confirm before generate: $confirm_gen" \
            "Show costs: $show_costs" \
            "Back")

        case "$choice" in
            "Auto-estimate:"*)
                [[ "$auto_estimate" == "true" ]] && auto_estimate="false" || auto_estimate="true"
                config_set preferences.auto_estimate "$auto_estimate"
                ;;
            "Confirm before generate:"*)
                [[ "$confirm_gen" == "true" ]] && confirm_gen="false" || confirm_gen="true"
                config_set preferences.confirm_before_generate "$confirm_gen"
                ;;
            "Show costs:"*)
                [[ "$show_costs" == "true" ]] && show_costs="false" || show_costs="true"
                config_set preferences.show_costs "$show_costs"
                ;;
            "Back"|"")
                break
                ;;
        esac
    done
}

# Cost tracking settings
settings_cost_tracking() {
    ui_info_box "Cost tracking helps you monitor API spending." "info"
    echo ""

    local budget
    budget=$(ui_input "Monthly budget (USD, 0 for unlimited)" "$(config_get cost_tracking.monthly_budget)")

    if [[ "$budget" =~ ^[0-9]+(\.[0-9]+)?$ ]]; then
        config_set cost_tracking.monthly_budget "$budget"
        config_set cost_tracking.enabled "true"
        success "Budget set to \$$budget USD per month"
    fi
}

# Security settings
settings_security() {
    local choice
    choice=$(ui_choose "Security Settings" \
        "Run security audit" \
        "Encrypt API keys" \
        "Test API keys" \
        "Back")

    case "$choice" in
        "Run security audit")
            security_audit
            ;;
        "Encrypt API keys")
            encrypt_keys
            ;;
        "Test API keys")
            test_all_keys
            ;;
    esac
}

# Test all configured API keys
test_all_keys() {
    msg "Testing configured API keys..."
    echo ""

    for provider in gemini claude openai groq xai; do
        if has_api_key "$provider"; then
            test_api_key "$provider"
        else
            debug "Skipping $provider (no key configured)"
        fi
    done

    echo ""
    success "Key testing complete"
}

# Reset cost counter
cmd_costreset() {
    local workspace logs_dir usage_log
    workspace="$(get_workspace)"
    logs_dir="$workspace/logs"
    usage_log="$logs_dir/usage.jsonl"

    # Get current total cost before resetting
    local current_cost
    current_cost=$(get_session_cost)

    if [[ ! -f "$usage_log" ]]; then
        info "No usage log found - cost is already $0.00"
        return 0
    fi

    # Show current cost
    echo ""
    warn "Current total cost: \$$current_cost"
    echo ""
    echo "This will:"
    echo "  • Delete the usage log file ($usage_log)"
    echo "  • Reset total cost to \$0.00"
    echo "  • Remove all cost history"
    echo ""

    if ! ui_confirm "Reset cost counter?" "no"; then
        msg "Cancelled"
        return 0
    fi

    # Backup the log before deleting (optional)
    local backup_file="$logs_dir/usage_backup_$(date +%Y%m%d_%H%M%S).jsonl"
    if ui_confirm "Create backup before resetting?" "yes"; then
        cp "$usage_log" "$backup_file"
        success "Backup created: $backup_file"
    fi

    # Remove the usage log
    rm -f "$usage_log"

    echo ""
    success "Cost counter reset to \$0.00"
    info "Previous cost was: \$$current_cost"
}

# Show session history - interactive menu
cmd_history() {
    local workspace logs_dir
    workspace="$(get_workspace)"
    logs_dir="$workspace/logs"

    if [[ ! -d "$logs_dir" ]]; then
        info "No chat history found"
        return 0
    fi

    while true; do
        local log_count=$(ls -1 "$logs_dir"/chat_*.log 2>/dev/null | wc -l | tr -d ' ')

        local choice
        choice=$(ui_choose "Chat History ($log_count sessions)" \
            "View recent sessions" \
            "Search in history" \
            "Show statistics" \
            "Export all history" \
            "Clear old logs" \
            "Back")

        case "$choice" in
            "View recent sessions")
                if [[ $log_count -eq 0 ]]; then
                    warn "No chat sessions found"
                    continue
                fi

                # List sessions with timestamps
                local logs sessions
                logs=($(ls -t "$logs_dir"/chat_*.log 2>/dev/null))
                sessions=()

                for log in "${logs[@]}"; do
                    local basename=$(basename "$log")
                    local date_part=$(echo "$basename" | sed 's/chat_\(.*\)\.log/\1/')
                    local formatted_date=$(echo "$date_part" | sed 's/_/ /g' | sed 's/\([0-9]\{4\}\)\([0-9]\{2\}\)\([0-9]\{2\}\) \([0-9]\{2\}\)\([0-9]\{2\}\)\([0-9]\{2\}\)/\1-\2-\3 \4:\5:\6/')
                    local msg_count=$(grep -c "USER:" "$log" 2>/dev/null || echo 0)
                    sessions+=("$formatted_date ($msg_count messages)")
                done
                sessions+=("Back")

                local selected
                selected=$(ui_choose "Select Session" "${sessions[@]}")

                if [[ -n "$selected" ]] && [[ "$selected" != "Back" ]]; then
                    # Find the corresponding log file
                    local idx=0
                    for i in "${!sessions[@]}"; do
                        if [[ "${sessions[$i]}" == "$selected" ]]; then
                            idx=$i
                            break
                        fi
                    done

                    if [[ $idx -lt ${#logs[@]} ]]; then
                        clear 2>/dev/null || true
                        ui_header "Chat Session: $selected"
                        echo ""
                        cat "${logs[$idx]}"
                        echo ""
                        ui_pause
                    fi
                fi
                ;;

            "Search in history")
                if [[ $log_count -eq 0 ]]; then
                    warn "No chat sessions to search"
                    continue
                fi

                local search_term
                search_term=$(ui_input "Search for:")
                if [[ -n "$search_term" ]]; then
                    echo ""
                    msg "Search results for: '$search_term'"
                    echo ""
                    grep -i -n --color=auto "$search_term" "$logs_dir"/chat_*.log 2>/dev/null || echo "No matches found"
                    echo ""
                    ui_pause
                fi
                ;;

            "Show statistics")
                if [[ $log_count -eq 0 ]]; then
                    warn "No chat sessions found"
                    continue
                fi

                clear 2>/dev/null || true
                ui_header "Chat History Statistics"
                echo ""

                local total_messages=$(grep -c "USER:" "$logs_dir"/chat_*.log 2>/dev/null || echo 0)
                local total_ai_responses=$(grep -c "AI:" "$logs_dir"/chat_*.log 2>/dev/null || echo 0)
                local oldest_log=$(ls -t "$logs_dir"/chat_*.log 2>/dev/null | tail -1)
                local newest_log=$(ls -t "$logs_dir"/chat_*.log 2>/dev/null | head -1)

                echo "Total sessions: $log_count"
                echo "Total user messages: $total_messages"
                echo "Total AI responses: $total_ai_responses"

                if [[ -n "$oldest_log" ]]; then
                    local oldest_date=$(basename "$oldest_log" | sed 's/chat_\(.*\)\.log/\1/' | sed 's/_/ /g')
                    echo "Oldest session: $oldest_date"
                fi

                if [[ -n "$newest_log" ]]; then
                    local newest_date=$(basename "$newest_log" | sed 's/chat_\(.*\)\.log/\1/' | sed 's/_/ /g')
                    echo "Latest session: $newest_date"
                fi

                # Calculate total size
                local total_size=$(du -sh "$logs_dir" 2>/dev/null | cut -f1)
                echo "Total storage: $total_size"

                echo ""
                ui_pause
                ;;

            "Export all history")
                if [[ $log_count -eq 0 ]]; then
                    warn "No chat sessions to export"
                    continue
                fi

                local export_file
                export_file=$(ui_input "Export to file:" "$HOME/aiwb_history_export.txt")
                if [[ -n "$export_file" ]]; then
                    cat "$logs_dir"/chat_*.log > "$export_file" 2>/dev/null
                    success "History exported to: $export_file"
                fi
                ;;

            "Clear old logs")
                if [[ $log_count -eq 0 ]]; then
                    info "No logs to clear"
                    continue
                fi

                local days
                days=$(ui_input "Delete logs older than (days):" "30")
                if [[ -n "$days" ]] && [[ "$days" =~ ^[0-9]+$ ]]; then
                    if ui_confirm "Delete logs older than $days days?" "no"; then
                        local deleted=0
                        find "$logs_dir" -name "chat_*.log" -type f -mtime +$days -exec rm {} \; -exec echo {} \; | while read -r; do
                            ((deleted++))
                        done
                        success "Deleted old chat logs"
                    fi
                fi
                ;;

            "Back"|"")
                break
                ;;
        esac
    done
}

# Show cost breakdown
cmd_costs() {
    local workspace logs_dir usage_log
    workspace="$(get_workspace)"
    logs_dir="$workspace/logs"
    usage_log="$logs_dir/usage.jsonl"

    if [[ ! -f "$usage_log" ]]; then
        info "No usage data found"
        return 0
    fi

    ui_header "Cost Breakdown (Estimated)"

    echo "${DIM}Note: Costs are estimated based on token approximation and may differ from actual API costs.${RESET}"
    echo "${DIM}Check your API provider's dashboard for actual usage and billing.${RESET}"
    echo ""

    # Total costs
    local total
    total=$(jq -s 'map(.cost) | add' "$usage_log" 2>/dev/null || echo "0")
    total=$(printf "%.2f" "$total")
    echo "Total estimated: \$${total} USD"
    echo ""

    # By provider
    echo "By Provider:"
    jq -s 'group_by(.provider) | map({provider: .[0].provider, cost: (map(.cost) | add)}) | .[]' "$usage_log" 2>/dev/null | \
        jq -r '"\(.provider)\t$\(.cost | tonumber | . * 100 | round / 100)"' | column -t

    echo ""

    # By date
    echo "Recent activity:"
    jq -s 'sort_by(.timestamp) | reverse | limit(10; .[]) | "\(.timestamp | split("T")[0])\t\(.provider)\t$\(.cost | tonumber | . * 100 | round / 100)"' "$usage_log" 2>/dev/null | column -t
}

# System health check
cmd_doctor() {
    ui_header "System Health Check"

    # Check dependencies
    echo "Dependencies:"
    for cmd in bash jq curl git gum fzf age; do
        if have "$cmd"; then
            success "✓ $cmd installed"
        else
            warn "○ $cmd not installed (optional)" 2>/dev/null || echo "○ $cmd not installed (optional)"
        fi
    done

    echo ""

    # Check API keys
    echo "API Keys:"
    for provider in gemini claude openai groq xai; do
        if has_api_key "$provider"; then
            success "✓ $provider key configured"
        else
            info "○ $provider key not configured"
        fi
    done

    # Check GitHub token
    echo ""
    echo "GitHub Integration:"
    if has_github_token; then
        success "✓ GitHub token configured"
    else
        info "○ GitHub token not configured (run: aiwb github auth)"
    fi

    echo ""

    # Check workspace
    local workspace
    workspace="$(get_workspace)"
    echo "Workspace: $workspace"
    if [[ -d "$workspace" ]]; then
        success "✓ Workspace exists"
    else
        warn "⚠ Workspace not found"
    fi

    echo ""

    # Platform info
    echo "Platform:"
    echo "  OS: $(get_platform)"
    echo "  Arch: $(uname -m)"
    echo "  Shell: $SHELL"

    echo ""
    success "Health check complete"
}

# Helper: copy the latest output file to the --output destination.
# Called both on success and on failure/timeout so partial output is preserved.
_headless_copy_output() {
    local output="$1"
    [[ -z "$output" ]] && return 0

    local workspace output_dir latest
    workspace="$(get_workspace)"
    output_dir="$workspace/outputs"
    latest=$(find "$output_dir" -name "*.md" ! -name "*.feedback.md" -type f 2>/dev/null | sort -r | head -1)
    if [[ -n "$latest" ]]; then
        cp "$latest" "$output"
        msg "Output copied to: $output"
        return 0
    else
        err "headless: no output file found to copy"
        return 1
    fi
}

# Headless (non-interactive) command handler for bot/CI use
cmd_headless() {
    local mode="" prompt="" instruct="" model="" provider="" model_name=""
    local verifier="" check_instruct="" context="" output=""
    local timeout_secs="${AIWB_HEADLESS_TIMEOUT}"

    # Parse flags
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --mode)       mode="$2";          shift 2 ;;
            --prompt)     prompt="$2";        shift 2 ;;
            --instruct)   instruct="$2";      shift 2 ;;
            --model)      model="$2";         shift 2 ;;
            --provider)   provider="$2";      shift 2 ;;
            --model-name) model_name="$2";    shift 2 ;;
            --verifier)   verifier="$2";      shift 2 ;;
            --check-instruct) check_instruct="$2"; shift 2 ;;
            --context)    context="$2";       shift 2 ;;
            --output)     output="$2";        shift 2 ;;
            --timeout)    timeout_secs="$2";  shift 2 ;;
            *) err "Unknown headless flag: $1"; return 1 ;;
        esac
    done

    # Validate required args
    if [[ -z "$mode" ]]; then
        err "headless: --mode is required (make|tweak|debug)"
        return 1
    fi
    if [[ -z "$prompt" && -z "$instruct" ]]; then
        err "headless: --prompt or --instruct is required"
        return 1
    fi

    # Set MODE_* state variables
    export AIWB_HEADLESS=1

    MODE_CURRENT=""
    MODE_PROMPT=""
    MODE_INSTRUCT_FILE=""
    MODE_MODEL_PROVIDER=""
    MODE_MODEL_NAME=""
    MODE_UPLOADS=()
    MODE_CHECK_PROVIDER=""
    MODE_CHECK_MODEL=""
    MODE_CHECK_INSTRUCT=""

    [[ -n "$provider" ]]       && MODE_MODEL_PROVIDER="$provider"
    [[ -n "$model_name" ]]     && MODE_MODEL_NAME="$model_name"
    # --model is shorthand for provider/name in "provider:name" form
    if [[ -n "$model" && "$model" == *:* ]]; then
        MODE_MODEL_PROVIDER="${model%%:*}"
        MODE_MODEL_NAME="${model##*:}"
    elif [[ -n "$model" ]]; then
        MODE_MODEL_NAME="$model"
    fi

    if [[ -n "$prompt" ]]; then
        MODE_PROMPT="$prompt"
    elif [[ -n "$instruct" ]]; then
        if [[ ! -f "$instruct" ]]; then
            err "headless: instruct file not found: $instruct"
            return 1
        fi
        MODE_INSTRUCT_FILE="$instruct"
    fi

    [[ -n "$check_instruct" ]] && MODE_CHECK_INSTRUCT="$check_instruct"
    if [[ -n "$verifier" && "$verifier" == *:* ]]; then
        MODE_CHECK_PROVIDER="${verifier%%:*}"
        MODE_CHECK_MODEL="${verifier##*:}"
    fi

    # Add context file/dir to uploads if provided
    if [[ -n "$context" ]]; then
        if [[ -e "$context" ]]; then
            mode_add_file "$context"
        else
            err "headless: context path not found: $context"
            return 1
        fi
    fi

    # --- Start watchdog timer for headless timeout ---
    local _watchdog_pid=""
    if [[ -n "$timeout_secs" ]] && [[ "$timeout_secs" -gt 0 ]] 2>/dev/null; then
        debug "headless: pipeline timeout set to ${timeout_secs}s"
        (
            sleep "$timeout_secs" 2>/dev/null
            # Time's up — try to save whatever output exists, then signal parent
            warn "headless: pipeline timed out after ${timeout_secs}s"
            kill -USR1 $$ 2>/dev/null || true
        ) &
        _watchdog_pid=$!
        # Trap USR1: on timeout, copy partial output then exit
        trap '
            warn "headless: timed out — saving partial output"
            _headless_copy_output "'"$output"'" 2>/dev/null || true
            kill '"$_watchdog_pid"' 2>/dev/null || true
            exit 124
        ' USR1
    fi

    # Run
    init_mode "$mode"
    local run_exit=0
    mode_run || run_exit=$?

    # Cancel watchdog if still running
    if [[ -n "$_watchdog_pid" ]]; then
        kill "$_watchdog_pid" 2>/dev/null || true
        wait "$_watchdog_pid" 2>/dev/null || true
    fi

    # Copy output to --output path if requested.
    # Always attempt this — even if mode_run failed (e.g. verifier timed out),
    # the generator output may already be saved. Partial success > no output.
    if [[ -n "$output" ]]; then
        if ! _headless_copy_output "$output"; then
            [[ $run_exit -eq 0 ]] && run_exit=1
        fi
    fi

    return $run_exit
}

# GitHub command handler
cmd_github() {
    local subcmd="${1:-}"
    shift || true

    case "$subcmd" in
        ""|menu)
            github_menu
            ;;
        status|stat)
            github_status
            ;;
        clone)
            github_clone "$@"
            ;;
        commit)
            github_commit "$@"
            ;;
        push)
            github_push "$@"
            ;;
        pull)
            github_pull "$@"
            ;;
        fetch)
            github_fetch "$@"
            ;;
        sync)
            github_sync "$@"
            ;;
        add|stage)
            github_add "$@"
            ;;
        branch|br)
            local branch_cmd="${1:-list}"
            shift || true
            case "$branch_cmd" in
                list)   github_branches "${1:-false}" ;;
                create) github_branch_create "$@" ;;
                switch) github_branch_switch "$@" ;;
                delete) github_branch_delete "$@" ;;
                *)      github_branches ;;
            esac
            ;;
        issue|issues)
            local issue_cmd="${1:-list}"
            shift || true
            case "$issue_cmd" in
                list)    github_issues_list "$@" ;;
                view)    github_issue_view "" "$@" ;;
                create)  github_issue_create "" "$@" ;;
                close)   github_issue_close "" "$@" ;;
                reopen)  github_issue_reopen "" "$@" ;;
                comment) github_issue_comment "" "$@" ;;
                *)       github_issues_list ;;
            esac
            ;;
        pr|prs)
            local pr_cmd="${1:-list}"
            shift || true
            case "$pr_cmd" in
                list)   github_pr_list "$@" ;;
                view)   github_pr_view "" "$@" ;;
                create) github_pr_create "" "$@" ;;
                merge)  github_pr_merge "" "$@" ;;
                close)  github_pr_close "" "$@" ;;
                files)  github_pr_files "" "$@" ;;
                *)      github_pr_list ;;
            esac
            ;;
        workflow|workflows|ci)
            local wf_cmd="${1:-list}"
            shift || true
            case "$wf_cmd" in
                list)  github_workflows_list "$@" ;;
                view)  github_workflow_view "" "$@" ;;
                rerun) github_workflow_rerun "" "$@" ;;
                *)     github_workflows_list ;;
            esac
            ;;
        auth|login)
            github_auth_setup
            ;;
        whoami|me)
            github_whoami
            ;;
        repo|info)
            github_repo_info "$@"
            ;;
        fork)
            github_fork "$@"
            ;;
        help|--help|-h)
            cmd_github_help
            ;;
        *)
            err "Unknown GitHub command: $subcmd"
            cmd_github_help
            return 1
            ;;
    esac
}

# GitHub help
cmd_github_help() {
    cat <<'EOF'
GitHub Integration Commands

USAGE:
  aiwb github [command] [options]

REPOSITORY:
  status              Show git status
  clone <repo>        Clone a repository
  commit <message>    Commit changes
  push [remote] [branch]   Push to remote
  pull [remote] [branch]   Pull from remote
  fetch [remote]      Fetch updates
  sync [remote]       Sync with remote (fetch, pull/push as needed)
  add [files...]      Stage files (or all if none specified)

BRANCHES:
  branch list         List branches
  branch create <name> [base]  Create new branch
  branch switch <name>         Switch branch
  branch delete <name>         Delete branch

ISSUES:
  issue list [state]  List issues (open/closed/all)
  issue view <num>    View issue details
  issue create <title> [body]  Create issue
  issue close <num>   Close issue
  issue reopen <num>  Reopen issue
  issue comment <num> <body>  Add comment

PULL REQUESTS:
  pr list [state]     List pull requests
  pr view <num>       View PR details
  pr create <title> [body] [base]  Create PR
  pr merge <num> [method]  Merge PR (merge/squash/rebase)
  pr close <num>      Close PR without merging
  pr files <num>      List changed files

WORKFLOWS (CI/CD):
  workflow list       List recent workflow runs
  workflow view <id>  View workflow run details
  workflow rerun <id> Re-run a workflow

USER:
  auth                Set up GitHub authentication
  whoami              Show authenticated user info
  repo [owner/repo]   Show repository info
  fork <owner/repo>   Fork a repository

EXAMPLES:
  aiwb github status
  aiwb github clone octocat/Hello-World
  aiwb github commit "Add new feature"
  aiwb github push
  aiwb github issue create "Bug report" "Description here"
  aiwb github pr create "New feature" "PR description" main
  aiwb github auth

For interactive mode: aiwb github (or aiwb gh)
EOF
}

# Repo context command
cmd_repo() {
    local subcmd="${1:-status}"

    case "$subcmd" in
        status|"")
            if is_repo_mode; then
                ui_header "Repository Context"
                echo "Repository: ${AIWB_REPO_NAME}"
                echo "Path:       ${AIWB_REPO_PATH}"
                echo "Branch:     ${AIWB_REPO_BRANCH}"
                echo "Remote:     ${AIWB_REPO_REMOTE:-none}"
                echo ""
                # Show git status
                if have git; then
                    echo "Status:"
                    git -C "$AIWB_REPO_PATH" status --short 2>/dev/null | head -10
                fi
            else
                info "No repository context active"
                echo ""
                echo "To enable repo mode:"
                echo "  1. Navigate to a git repository: cd /path/to/repo"
                echo "  2. Start AIWB: aiwb"
                echo ""
                echo "Or clone a repo: aiwb github clone owner/repo"
            fi
            ;;
        none|disable|off)
            AIWB_REPO_ENABLED=false
            AIWB_REPO_PATH=""
            AIWB_REPO_NAME=""
            success "Repository context disabled"
            ;;
        refresh|detect)
            detect_repo
            if is_repo_mode; then
                success "Detected repo: $AIWB_REPO_NAME ($AIWB_REPO_BRANCH)"
            else
                info "No git repository found in current directory"
            fi
            ;;
        *)
            echo "Usage: /repo [status|none|refresh]"
            echo ""
            echo "Commands:"
            echo "  status   Show current repo context (default)"
            echo "  none     Disable repo mode"
            echo "  refresh  Re-detect repository"
            ;;
    esac
}

# ============================================================================
# MAIN
# ============================================================================

main() {
    # Parse global flags
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --debug)
                export AIWB_DEBUG=1
                shift
                ;;
            --version|-v)
                echo "AIWB v3.0.0"
                exit 0
                ;;
            --help|-h)
                cmd_help
                exit 0
                ;;
            --provider)
                config_set model_provider "$2"
                shift 2
                ;;
            --model)
                config_set model_name "$2"
                shift 2
                ;;
            *)
                break
                ;;
        esac
    done

    # Initialize
    initialize

    # Route command
    local cmd="${1:-chat}"
    shift || true

    debug "Routing to command: $cmd"

    case "$cmd" in
        chat|c)
            cmd_chat "$@"
            ;;
        quick|q)
            cmd_quick "$@"
            ;;
        wizard|wiz)
            cmd_wizard "$@"
            ;;
        make|m)
            cmd_make "$@"
            ;;
        headless)
            cmd_headless "$@"
            ;;
        improve|imp)
            cmd_improve "$@"
            ;;
        templates|tmpl)
            cmd_templates "$@"
            ;;
        generate|gen|g)
            cmd_generate "$@"
            ;;
        estimate|est|e)
            cmd_estimate "$@"
            ;;
        verify|review)
            cmd_verify "$@"
            ;;
        refine|loop)
            cmd_refine "$@"
            ;;
        keys|key)
            cmd_keys "$@"
            ;;
        models|settings|config)
            cmd_models "$@"
            ;;
        status|stat)
            ui_show_status
            ;;
        history|hist)
            cmd_history "$@"
            ;;
        costs|cost)
            cmd_costs "$@"
            ;;
        doctor|check)
            cmd_doctor "$@"
            ;;
        security-audit|audit)
            security_audit "$@"
            ;;
        scanrepo|scan)
            cmd_scanrepo "$@"
            ;;
        smartscan|smart)
            cmd_smartscan "$@"
            ;;
        github|gh)
            cmd_github "$@"
            ;;
        help|--help|-h)
            cmd_help
            ;;        
              *)
            # If no command is found, default to chat
            cmd_chat "$@"
            ;;




    esac
}

# Run main
main "$@"