#!/bin/bash
# costs.sh - Cost tracking and budget management for AI Consultants
#
# Tracks consultation costs based on estimated token usage
# and manages budget limits.

# =============================================================================
# EXTERNAL RATES FILE (v2.4)
# =============================================================================
# Path to external rates file for easy updates

# Get the script directory for relative path resolution
_COSTS_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
COST_RATES_FILE="${COST_RATES_FILE:-$_COSTS_SCRIPT_DIR/../../docs/cost_rates.json}"

# Load rate from JSON file
# Usage: get_rate_from_file <model> <type: input|output>
# Returns: rate as string, or exits with 1 if not found
#
# The match is CASE-INSENSITIVE: callers lowercase the model name, but several
# rate keys are mixed-case (agy display names like "Gemini 3.1 Pro (High)", the
# display-name model IDs). An exact-case jq lookup on a
# lowercased model would miss those keys and silently fall through to
# default_rate — mis-billing free/local and Gemini models. Downcasing both sides
# fixes it generally (lowercase keys like minimax-m2.7/gpt-5.5 still match).
get_rate_from_file() {
    local model="$1"
    local type="$2"

    if [[ -f "$COST_RATES_FILE" ]]; then
        local rate
        rate=$(jq -r --arg m "$model" --arg t "$type" \
            'first(.models | to_entries[] | select((.key | ascii_downcase) == ($m | ascii_downcase)) | .value[$t]) // null' \
            "$COST_RATES_FILE" 2>/dev/null)
        if [[ "$rate" != "null" && -n "$rate" ]]; then
            echo "$rate"
            return 0
        fi
    fi
    return 1
}

# Get fallback model for a consultant
# Usage: get_consultant_fallback_model <consultant>
# Returns: fallback model name (empty string if not found)
get_consultant_fallback_model() {
    local consultant="$1"
    consultant=$(echo "$consultant" | tr '[:upper:]' '[:lower:]')

    if [[ -f "$COST_RATES_FILE" ]]; then
        local fallback
        fallback=$(jq -r ".consultant_fallbacks[\"$consultant\"] // null" "$COST_RATES_FILE" 2>/dev/null)
        if [[ "$fallback" != "null" && -n "$fallback" ]]; then
            echo "$fallback"
        fi
    fi
}

# Resolve model name: use reported model if known, else fallback
# Usage: resolve_model_for_cost <reported_model> <consultant>
# Returns: resolved model name for cost calculation
resolve_model_for_cost() {
    local model="$1"
    local consultant="$2"

    # Fetch fallback once for reuse
    local fallback
    fallback=$(get_consultant_fallback_model "$consultant")

    # If model is "default", empty, or unknown, use fallback
    if [[ -z "$model" || "$model" == "default" ]]; then
        if [[ -n "$fallback" ]]; then
            type log_debug &>/dev/null && log_debug "Using fallback model '$fallback' for $consultant (reported: '$model')"
            echo "$fallback"
            return 0
        fi
    fi

    # Check if model exists in rates file
    if get_rate_from_file "$model" "input" >/dev/null 2>&1; then
        echo "$model"
        return 0
    fi

    # Model not in rates file, try fallback
    if [[ -n "$fallback" ]]; then
        type log_debug &>/dev/null && log_debug "Model '$model' not in rates, using fallback '$fallback' for $consultant"
        echo "$fallback"
        return 0
    fi

    # Return original (will use default rate)
    echo "$model"
}

# Get default rate from JSON file
# Usage: get_default_rate <type: input|output>
get_default_rate() {
    local type="$1"
    if [[ -f "$COST_RATES_FILE" ]]; then
        local rate
        rate=$(jq -r ".default_rate.$type // null" "$COST_RATES_FILE" 2>/dev/null)
        if [[ "$rate" != "null" && -n "$rate" ]]; then
            echo "$rate"
            return 0
        fi
    fi
    # Hardcoded fallback
    case "$type" in
        input)  echo "0.005" ;;
        output) echo "0.015" ;;
    esac
}

# =============================================================================
# COST RATES (USD per 1K tokens)
# =============================================================================
# Using case statements for bash 3.2 compatibility (no associative arrays)
# External JSON file is tried first, then fallback to hardcoded rates

# Get input token cost per 1K tokens
# Usage: get_input_cost_per_1k <model>
get_input_cost_per_1k() {
    local model="$1"
    model=$(echo "$model" | tr '[:upper:]' '[:lower:]')

    # Try external file first
    local rate
    if rate=$(get_rate_from_file "$model" "input" 2>/dev/null); then
        echo "$rate"
        return
    fi

    # Fallback to hardcoded rates for backwards compatibility
    case "$model" in
        gemini-2.5-pro)   echo "0.00125" ;;
        gemini-2.5-flash) echo "0.000075" ;;
        gemini-2.0-flash) echo "0.0001" ;;
        deepseek-flash|deepseek-v4-flash|deepseek-v4-flash-vision-exp) echo "0.0003" ;;
        gpt-6-astra|claude-fable-5-1) echo "0.01" ;;
        gpt-4)            echo "0.03" ;;
        gpt-4-turbo)      echo "0.01" ;;
        gpt-4o)           echo "0.005" ;;
        gpt-4o-mini)      echo "0.00015" ;;
        o1)               echo "0.015" ;;
        o3)               echo "0.015" ;;
        claude-3-opus)    echo "0.015" ;;
        claude-3-sonnet)  echo "0.003" ;;
        claude-3-haiku)   echo "0.00025" ;;
        mistral-large)    echo "0.004" ;;
        mistral-medium)   echo "0.0027" ;;
        mistral-small)    echo "0.001" ;;
        cursor)           echo "0.005" ;;
        # Qwen3 models (Alibaba DashScope)
        qwen-max)         echo "0.004" ;;
        qwen-plus)        echo "0.002" ;;
        qwen-turbo)       echo "0.0008" ;;
        # GLM models (Zhipu AI)
        glm-4)            echo "0.003" ;;
        glm-3-turbo)      echo "0.001" ;;
        # Grok models (xAI)
        grok-4.20-0309-reasoning)        echo "0.005" ;;
        grok-2)           echo "0.01" ;;
        # Default
        *)                echo "0.005" ;;
    esac
}

# Get output token cost per 1K tokens
# Usage: get_output_cost_per_1k <model>
get_output_cost_per_1k() {
    local model="$1"
    model=$(echo "$model" | tr '[:upper:]' '[:lower:]')

    # Try external file first
    local rate
    if rate=$(get_rate_from_file "$model" "output" 2>/dev/null); then
        echo "$rate"
        return
    fi

    # Fallback to hardcoded rates for backwards compatibility
    case "$model" in
        gemini-2.5-pro)   echo "0.005" ;;
        gemini-2.5-flash) echo "0.0003" ;;
        gemini-2.0-flash) echo "0.0004" ;;
        deepseek-flash|deepseek-v4-flash|deepseek-v4-flash-vision-exp) echo "0.0012" ;;
        gpt-6-astra|claude-fable-5-1) echo "0.05" ;;
        gpt-4)            echo "0.06" ;;
        gpt-4-turbo)      echo "0.03" ;;
        gpt-4o)           echo "0.015" ;;
        gpt-4o-mini)      echo "0.0006" ;;
        o1)               echo "0.06" ;;
        o3)               echo "0.06" ;;
        claude-3-opus)    echo "0.075" ;;
        claude-3-sonnet)  echo "0.015" ;;
        claude-3-haiku)   echo "0.00125" ;;
        mistral-large)    echo "0.012" ;;
        mistral-medium)   echo "0.0081" ;;
        mistral-small)    echo "0.003" ;;
        cursor)           echo "0.015" ;;
        # Qwen3 models (Alibaba DashScope)
        qwen-max)         echo "0.012" ;;
        qwen-plus)        echo "0.006" ;;
        qwen-turbo)       echo "0.002" ;;
        # GLM models (Zhipu AI)
        glm-4)            echo "0.009" ;;
        glm-3-turbo)      echo "0.003" ;;
        # Grok models (xAI)
        grok-4.20-0309-reasoning)        echo "0.015" ;;
        grok-2)           echo "0.03" ;;
        # Default
        *)                echo "0.015" ;;
    esac
}

# =============================================================================
# COST CALCULATION
# =============================================================================

# Estimate cost for a query
# Usage: estimate_query_cost <model> <input_tokens> <output_tokens>
estimate_query_cost() {
    local model="$1"
    local input_tokens="${2:-1000}"
    local output_tokens="${3:-500}"

    # Normalize model name
    model=$(echo "$model" | tr '[:upper:]' '[:lower:]')

    # Get rates using lookup functions
    local input_rate output_rate
    input_rate=$(get_input_cost_per_1k "$model")
    output_rate=$(get_output_cost_per_1k "$model")

    # Long-context multiplier applies to the entire request, not only excess.
    if [[ "$model" == gpt-6-astra && "$input_tokens" -gt 272000 ]]; then
        input_rate=$(echo "scale=6; $input_rate * 2" | bc)
        output_rate=$(echo "scale=6; $output_rate * 1.5" | bc)
    fi
    # Calculate cost
    local input_cost output_cost total_cost
    input_cost=$(echo "scale=6; $input_tokens / 1000 * $input_rate" | bc)
    output_cost=$(echo "scale=6; $output_tokens / 1000 * $output_rate" | bc)
    total_cost=$(echo "scale=6; $input_cost + $output_cost" | bc)

    # bc prints sub-1 values without a leading zero (".014000"); restore it so
    # cost reports read "0.014000". Numeric comparisons are unaffected either way.
    [[ "$total_cost" == .* ]] && total_cost="0$total_cost"
    [[ "$total_cost" == -.* ]] && total_cost="-0${total_cost#-}"

    echo "$total_cost"
}

# Check whether a model is billed by prepaid credits rather than per token.
# Usage: is_unpriced_model <model>
# Returns 0 if the model has no meaningful per-token price, 1 otherwise.
is_unpriced_model() {
    local model="$1"
    [[ -z "$model" ]] && return 1
    [[ -f "$COST_RATES_FILE" ]] || return 1

    local hit
    hit=$(jq -r --arg m "$model" \
        'first((.unpriced_models // [])[] | select(ascii_downcase == ($m | ascii_downcase))) // ""' \
        "$COST_RATES_FILE" 2>/dev/null)

    [[ -n "$hit" ]]
}

# Enumerate the response files a session actually billed for.
# Usage: _billable_response_files <responses_dir>
#
# Emits one path per line. Three rules, each of which was a real mis-billing
# before v2.25.0 made tokens_used non-zero and turned them into money:
#   - only consultant responses (voting.json and friends carry no tokens_used
#     and would each be priced at the `// 1000` fallback times the default rate);
#   - debate rounds live in `round_N/` subdirectories and are separate billed
#     queries, so those directories must be included without recursing into
#     derived peer-review trees;
#   - a cache hit made no API call, and `<consultant>_escalated.json` is a copy
#     of a file already counted, so both are excluded.
_billable_response_files() {
    local responses_dir="$1"
    [[ -d "$responses_dir" ]] || return 0

    local f round_dir round_name
    for f in "$responses_dir"/*.json "$responses_dir"/round_*/*.json; do
        [[ -f "$f" && -s "$f" ]] || continue

        # Nested billable files are allowed only one level below an exact
        # round_<digits> directory. Peer-review and other derived trees are not
        # provider queries and must never reach the token fallback.
        if [[ "${f%/*}" != "$responses_dir" ]]; then
            round_dir="${f%/*}"
            round_name="${round_dir##*/}"
            case "$round_name" in
                round_[0-9]*)
                    case "${round_name#round_}" in
                        *[!0-9]*) continue ;;
                    esac
                    ;;
                *) continue ;;
            esac
        fi

        [[ "${f##*/}" == *_escalated.json ]] && continue
        _is_consultant_response_file "$f" 2>/dev/null || continue
        # A cache hit is a replay of an earlier response, metadata included.
        [[ "$(jq -r '.cache_metadata.from_cache // false' "$f" 2>/dev/null)" == "true" ]] && continue
        printf '%s\n' "$f"
    done
}

# Resolve the catalog key used for billing without relabeling the response's
# effective identity. A provider-reported dated revision may be absent from the
# catalog; in that case the explicit requested_model is the only defensible
# pricing key. Consultant premium fallbacks are used only for legacy responses
# that carry no requested_model at all.
_billing_model_for_response() {
    local f="$1" effective requested consultant
    effective=$(jq -r '.model // "default"' "$f" 2>/dev/null)
    requested=$(jq -r '.metadata.requested_model // empty' "$f" 2>/dev/null)
    consultant=$(jq -r '.consultant // empty' "$f" 2>/dev/null)

    if get_rate_from_file "$effective" input >/dev/null 2>&1 \
        || is_unpriced_model "$effective"; then
        printf '%s\n' "$effective"
    elif [[ -n "$requested" ]]; then
        printf '%s\n' "$requested"
    else
        resolve_model_for_cost "$effective" "$consultant"
    fi
}

# Calculate total session cost from responses
# Usage: calculate_session_cost <responses_dir>
#
# Prefers an exact provider cost when recorded (Claude CLI), then the provider's
# prompt/completion split (API mode), and only then the legacy 60/40 guess.
calculate_session_cost() {
    local responses_dir="$1"
    local total_cost=0
    local f

    while IFS= read -r f; do
        [[ -n "$f" ]] || continue
        local provider_cost
        provider_cost=$(jq -r '.metadata.provider_cost_usd // empty' "$f" 2>/dev/null)
        if [[ "$provider_cost" =~ ^[0-9]+([.][0-9]+)?$ ]]; then
            total_cost=$(echo "scale=6; $total_cost + $provider_cost" | bc)
            continue
        fi

        local model input_tokens output_tokens
        model=$(_billing_model_for_response "$f")
        input_tokens=$(jq -r '.metadata.tokens_input // empty' "$f" 2>/dev/null)
        output_tokens=$(jq -r '.metadata.tokens_output // empty' "$f" 2>/dev/null)

        if [[ -z "$input_tokens" || -z "$output_tokens" ]]; then
            local tokens
            tokens=$(jq -r '.metadata.tokens_used // 1000' "$f" 2>/dev/null)
            [[ "$tokens" =~ ^[0-9]+$ ]] || tokens=0
            input_tokens=$((tokens * 60 / 100))
            output_tokens=$((tokens * 40 / 100))
        fi
        [[ "$input_tokens"  =~ ^[0-9]+$ ]] || input_tokens=0
        [[ "$output_tokens" =~ ^[0-9]+$ ]] || output_tokens=0

        local cost
        cost=$(estimate_query_cost "$model" "$input_tokens" "$output_tokens")
        total_cost=$(echo "scale=6; $total_cost + $cost" | bc)
    done < <(_billable_response_files "$responses_dir")

    # Provider-reported costs may carry fewer decimal places than locally
    # estimated values. Normalize the public result to the existing 6-place
    # contract without changing its numeric value.
    total_cost=$(echo "scale=6; $total_cost / 1" | bc)

    # bc drops the leading zero on sub-1 values (".115000"); restore it as
    # estimate_query_cost does, so callers can compare numerically.
    [[ "$total_cost" == .* ]] && total_cost="0$total_cost"

    echo "$total_cost"
}

# Render the caveats a session's cost figure carries, or nothing when it is a
# clean measurement of priced models.
# Usage: format_cost_caveats <responses_dir>
#
# Two things can make the number less than it looks:
#   - credit-billed models contribute 0 because they have no per-token price;
#   - CLI-mode consultants without provider usage have a 4-chars-per-token
#     approximation (metadata.tokens_source = "estimated"). Claude CLI is the
#     exception: its JSON envelope supplies measured usage and exact cost.
# Both are stated rather than folded silently into a confident-looking total.
#
# Rescans the directory rather than reading state left by
# calculate_session_cost: that function is always invoked in a command
# substitution, so anything it assigns dies with the subshell and the
# disclosure would never fire.
format_cost_caveats() {
    local responses_dir="${1:-}"
    [[ -d "$responses_dir" ]] || return 0

    local estimated=0 unknown=0 priced=0 f astra_estimate=0 deepseek_flash_estimate=0
    local unpriced_models=()
    while IFS= read -r f; do
        [[ -n "$f" ]] || continue

        local src
        src=$(jq -r '.metadata.tokens_source // "unknown"' "$f" 2>/dev/null)
        case "$src" in
            estimated) estimated=$((estimated + 1)); priced=$((priced + 1)) ;;
            # Anything without a source contributed no tokens at all - in
            # practice a failed consultation. Counting it in the denominator of
            # "estimated for N of M" would make it read as provider-measured,
            # the exact opposite of the truth.
            measured)  priced=$((priced + 1)) ;;
            *)         unknown=$((unknown + 1)) ;;
        esac

        local model seen item
        model=$(_billing_model_for_response "$f")
        if [[ "$model" == gpt-6-astra ]] && ! jq -e '.metadata.provider_cost_usd | numbers' "$f" >/dev/null 2>&1; then
            astra_estimate=1
        fi
        case "$model" in
            deepseek-flash|deepseek-v4-flash|deepseek-v4-flash-vision-exp)
                if ! jq -e '.metadata.provider_cost_usd | numbers' "$f" >/dev/null 2>&1; then
                    deepseek_flash_estimate=1
                fi
                ;;
        esac
        is_unpriced_model "$model" || continue
        seen=false
        for item in "${unpriced_models[@]+"${unpriced_models[@]}"}"; do
            [[ "$item" == "$model" ]] && seen=true
        done
        [[ "$seen" == "true" ]] || unpriced_models+=("$model")
    done < <(_billable_response_files "$responses_dir")

    local parts=""
    if [[ $estimated -gt 0 ]]; then
        parts="token counts estimated for $estimated of $priced"
    fi
    if [[ $unknown -gt 0 ]]; then
        parts="${parts:+$parts; }$unknown contributed no token data"
    fi
    if [[ $astra_estimate -eq 1 ]]; then
        parts="${parts:+$parts; }cost estimated using Astra Standard rates; cache/service adjustments excluded, not an invoice"
    fi
    if [[ $deepseek_flash_estimate -eq 1 ]]; then
        parts="${parts:+$parts; }DeepSeek Flash cost estimated at peak cache-miss rates; off-peak/cache discounts excluded, not an invoice"
    fi
    local unpriced=""
    if (( ${#unpriced_models[@]} > 0 )); then
        unpriced=$(IFS=,; printf '%s' "${unpriced_models[*]}")
        parts="${parts:+$parts; }excludes ${unpriced//,/, } (credit-billed, no per-token price)"
    fi

    [[ -z "$parts" ]] && return 0
    echo " [$parts]"
}

# Format cost for display
# Usage: format_cost <cost_usd>
format_cost() {
    local cost="$1"

    # Convert to cents if very small
    if (( $(echo "$cost < 0.01" | bc -l) )); then
        local cents=$(echo "scale=2; $cost * 100" | bc)
        # bc drops the leading zero on sub-1 values (".03"); restore it.
        [[ "$cents" == .* ]] && cents="0$cents"
        echo "${cents}¢"
    else
        printf "\$%.4f" "$cost"
    fi
}

# =============================================================================
# BUDGET MANAGEMENT
# =============================================================================

# Check if cost exceeds budget
# Usage: check_budget <cost> <budget>
check_budget() {
    local cost="$1"
    local budget="${2:-${MAX_SESSION_COST:-1.00}}"

    if (( $(echo "$cost > $budget" | bc -l) )); then
        return 1  # Over budget
    fi
    return 0  # Within budget
}

# Check if we are close to warning threshold
# Usage: check_warning_threshold <cost>
check_warning_threshold() {
    local cost="$1"
    local threshold="${WARN_AT_COST:-0.50}"

    if (( $(echo "$cost > $threshold" | bc -l) )); then
        return 0  # Should warn
    fi
    return 1  # No warning needed
}

# Estimate pre-consultation cost (before executing)
# Usage: estimate_consultation_cost <num_consultants> <context_size_chars> [consultants_csv]
# The consultants_csv parameter is a comma-separated list of consultant names (e.g., "Gemini,Codex,Mistral")
estimate_consultation_cost() {
    local num_consultants="${1:-5}"
    local context_size="${2:-5000}"
    local consultants="${3:-}"

    # Estimate tokens from context (approximately 4 chars per token)
    local estimated_input_tokens=$((context_size / 4))

    # Estimate output tokens (approximately 500-1000 per response)
    local estimated_output_tokens=750

    local total=0

    if [[ -n "$consultants" ]]; then
        # Use provided consultant list with fallback models
        local IFS=','
        read -ra consultant_list <<< "$consultants"
        for consultant in "${consultant_list[@]}"; do
            local fallback_model
            fallback_model=$(get_consultant_fallback_model "$consultant")
            local model_to_use="${fallback_model:-default}"
            local cost=$(estimate_query_cost "$model_to_use" "$estimated_input_tokens" "$estimated_output_tokens")
            total=$(echo "scale=6; $total + $cost" | bc)
        done
    else
        # Fallback to generic estimate using default rate
        for ((i=0; i<num_consultants; i++)); do
            local default_input default_output
            default_input=$(get_default_rate "input")
            default_output=$(get_default_rate "output")
            local input_cost output_cost cost
            input_cost=$(echo "scale=6; $estimated_input_tokens / 1000 * $default_input" | bc)
            output_cost=$(echo "scale=6; $estimated_output_tokens / 1000 * $default_output" | bc)
            cost=$(echo "scale=6; $input_cost + $output_cost" | bc)
            total=$(echo "scale=6; $total + $cost" | bc)
        done
    fi

    echo "$total"
}

# =============================================================================
# COST TRACKING
# =============================================================================

# File for cumulative tracking (XDG-aware fallback added in v2.13)
COST_TRACKING_FILE="${COST_TRACKING_FILE:-${_AI_CONSULTANTS_XDG_DATA:-/tmp/ai_consultants}/costs.json}"

# Record session cost
# Usage: track_session_cost <session_id> <cost>
# Best-effort: cost tracking is bookkeeping. The caller in consult_all.sh runs
# under set -e *after* every consultant has already been queried and billed,
# so failures here degrade to a warning instead of aborting the run.
track_session_cost() {
    local session_id="$1"
    local cost="$2"

    # The XDG data dir may not exist yet on a fresh install — costs.json is
    # the only artifact stored there
    if ! mkdir -p "$(dirname "$COST_TRACKING_FILE")" 2>/dev/null; then
        log_warn "Cost tracking skipped: cannot create $(dirname "$COST_TRACKING_FILE")"
        return 0
    fi

    # Serialize the read-modify-write against concurrent consultations with a
    # portable mkdir lock (flock is unavailable on macOS). Bounded wait, then
    # proceed unlocked: bookkeeping must never block or abort the run.
    local lock_dir="${COST_TRACKING_FILE}.lock" locked=false _i
    for _i in {1..50}; do
        if mkdir "$lock_dir" 2>/dev/null; then locked=true; break; fi
        sleep 0.1
    done
    if [[ "$locked" != "true" ]]; then
        log_warn "Cost tracking lock busy for 5s (stale ${lock_dir}?), proceeding unlocked"
    fi

    _track_session_cost_update "$session_id" "$cost" || true

    if [[ "$locked" == "true" ]]; then
        rmdir "$lock_dir" 2>/dev/null || true
    fi
    return 0
}

# Inner update for track_session_cost — runs with the lock held.
# Never propagates failure; logs a warning and returns instead.
_track_session_cost_update() {
    local session_id="$1"
    local cost="$2"
    local timestamp=$(date -Iseconds)

    # A corrupt file (truncated write, interleaved concurrent update) would
    # fail every future jq update and never self-heal: set it aside and reset
    if [[ -f "$COST_TRACKING_FILE" ]]; then
        if ! jq empty "$COST_TRACKING_FILE" 2>/dev/null; then
            log_warn "Cost tracking file corrupt, resetting (backup: ${COST_TRACKING_FILE}.corrupt)"
            mv -f "$COST_TRACKING_FILE" "${COST_TRACKING_FILE}.corrupt" 2>/dev/null || true
        fi
    fi
    if [[ ! -f "$COST_TRACKING_FILE" ]]; then
        if ! { echo '{"sessions": [], "total_cost": 0}' > "$COST_TRACKING_FILE"; } 2>/dev/null; then
            log_warn "Cost tracking skipped: cannot write $COST_TRACKING_FILE"
            return 0
        fi
    fi

    # Unique temp file: even an unlocked writer must not share a fixed .tmp
    # sibling with other runs (lost records, failed mv)
    local tmp_file
    if ! tmp_file=$(mktemp "${COST_TRACKING_FILE}.XXXXXX" 2>/dev/null); then
        log_warn "Cost tracking skipped: cannot create temp file for $COST_TRACKING_FILE"
        return 0
    fi
    if jq --arg id "$session_id" \
          --arg cost "$cost" \
          --arg ts "$timestamp" \
          '.sessions += [{id: $id, cost: ($cost | tonumber), timestamp: $ts}] | .total_cost += ($cost | tonumber)' \
          "$COST_TRACKING_FILE" > "$tmp_file" 2>/dev/null; then
        mv -f "$tmp_file" "$COST_TRACKING_FILE" 2>/dev/null || rm -f "$tmp_file"
    else
        rm -f "$tmp_file"
        log_warn "Cost tracking update failed for session $session_id"
    fi
    return 0
}

# Get total tracked cost
# Usage: get_total_tracked_cost
get_total_tracked_cost() {
    if [[ -f "$COST_TRACKING_FILE" ]]; then
        jq -r '.total_cost // 0' "$COST_TRACKING_FILE"
    else
        echo 0
    fi
}

# Generate cost report
# Usage: generate_cost_report
generate_cost_report() {
    if [[ ! -f "$COST_TRACKING_FILE" ]]; then
        echo "No cost data available"
        return
    fi

    local total=$(jq -r '.total_cost // 0' "$COST_TRACKING_FILE")
    local sessions=$(jq -r '.sessions | length' "$COST_TRACKING_FILE")
    local avg=0
    if [[ $sessions -gt 0 ]]; then
        avg=$(echo "scale=6; $total / $sessions" | bc)
    fi

    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║                    Cost Report                               ║"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo ""
    echo "  Total sessions: $sessions"
    echo "  Total cost: $(format_cost $total)"
    echo "  Average per session: $(format_cost $avg)"
    echo ""

    # Last 5 sessions
    echo "  Recent sessions:"
    jq -r '.sessions | .[-5:] | .[] | "    \(.timestamp): \(.id) - $\(.cost)"' "$COST_TRACKING_FILE" 2>/dev/null || echo "    No sessions"
}

# =============================================================================
# RESPONSE LENGTH LIMITS (v2.3)
# =============================================================================

# Get max response tokens for a category
# Usage: get_max_response_tokens <category>
get_max_response_tokens() {
    local category="$1"
    local limits="${MAX_RESPONSE_TOKENS_BY_CATEGORY:-QUICK_SYNTAX:200,CODE_REVIEW:800,BUG_DEBUG:800,ARCHITECTURE:1000,SECURITY:1000,DATABASE:600,GENERAL:500}"

    # Search for category in the limits string
    local limit
    limit=$(echo "$limits" | tr ',' '\n' | grep -i "^${category}:" | cut -d: -f2 | head -1)

    # Return limit or default
    if [[ -n "$limit" && "$limit" =~ ^[0-9]+$ ]]; then
        echo "$limit"
    else
        echo "500"  # Default
    fi
}

# Check if response limits are enabled
# Usage: is_response_limits_enabled
# NOTE: Default is false (opt-in) per quality review - aligns with config.sh
is_response_limits_enabled() {
    [[ "${ENABLE_RESPONSE_LIMITS:-false}" == "true" ]]
}

# Get model tier (economy, standard, premium)
# Usage: get_model_tier <model>
get_model_tier() {
    local model="$1"
    model=$(echo "$model" | tr '[:upper:]' '[:lower:]')

    case "$model" in
        # Economy tier - cheapest models
        gemini-2.5-flash|gemini-2.0-flash|gpt-4o-mini|claude-3-haiku|mistral-small|qwen-turbo|glm-3-turbo)
            echo "economy"
            ;;
        # Premium tier - most expensive models
        gpt-4|gpt-4-turbo|o1|o3|claude-3-opus|mistral-large|grok-2|qwen-max)
            echo "premium"
            ;;
        # Standard tier - default/mid-range
        *)
            echo "standard"
            ;;
    esac
}

# Get economic model for a consultant
# Delegates to get_model_for_tier() in config.sh (single source of truth)
# Usage: get_economic_model <consultant> [cli|api]
get_economic_model() {
    local consultant="$1"
    local transport="${2:-}"
    if type get_model_for_tier &>/dev/null; then
        get_model_for_tier "$consultant" "economy" "$transport"
    else
        echo ""
    fi
}

# Calculate query complexity score (1-10)
# Usage: calculate_query_complexity <query> <num_files> <category>
calculate_query_complexity() {
    local query="$1"
    local num_files="${2:-0}"
    local category="${3:-GENERAL}"

    local score=5  # Base score

    # Length factor
    local query_len=${#query}
    if [[ $query_len -gt 500 ]]; then
        score=$((score + 2))
    elif [[ $query_len -gt 200 ]]; then
        score=$((score + 1))
    elif [[ $query_len -lt 50 ]]; then
        score=$((score - 1))
    fi

    # File count factor
    if [[ $num_files -gt 5 ]]; then
        score=$((score + 2))
    elif [[ $num_files -gt 2 ]]; then
        score=$((score + 1))
    elif [[ $num_files -eq 0 ]]; then
        score=$((score - 1))
    fi

    # Category factor
    case "$category" in
        ARCHITECTURE|SECURITY)
            score=$((score + 2))
            ;;
        CODE_REVIEW|BUG_DEBUG)
            score=$((score + 1))
            ;;
        QUICK_SYNTAX)
            score=$((score - 2))
            ;;
    esac

    # Keyword complexity indicators
    if echo "$query" | grep -qiE "(architecture|design|scalab|security|performance|refactor|migrate)"; then
        score=$((score + 1))
    fi
    if echo "$query" | grep -qiE "(fix|bug|error|typo|rename)"; then
        score=$((score - 1))
    fi

    # Cap at 1-10
    [[ $score -gt 10 ]] && score=10
    [[ $score -lt 1 ]] && score=1

    echo "$score"
}

# Check if query is simple (should use economic models)
# Usage: is_simple_query <complexity_score>
is_simple_query() {
    local complexity="${1:-5}"
    local threshold="${COMPLEXITY_THRESHOLD_SIMPLE:-3}"

    [[ $complexity -le $threshold ]]
}

# Check if query is complex (should use premium models)
# Usage: is_complex_query <complexity_score>
is_complex_query() {
    local complexity="${1:-5}"
    local threshold="${COMPLEXITY_THRESHOLD_MEDIUM:-6}"

    [[ $complexity -gt $threshold ]]
}

# =============================================================================
# BUDGET ENFORCEMENT (v2.4)
# =============================================================================

# Check if budget enforcement is enabled
# Usage: is_budget_enabled
is_budget_enabled() {
    [[ "${ENABLE_BUDGET_LIMIT:-false}" == "true" ]]
}

# Get remaining budget
# Usage: get_remaining_budget <current_cost>
get_remaining_budget() {
    local current_cost="${1:-0}"
    local budget="${MAX_SESSION_COST:-1.00}"

    echo "scale=6; $budget - $current_cost" | bc
}

# Format budget status for display
# Usage: format_budget_status <current_cost> [budget]
format_budget_status() {
    local current_cost="${1:-0}"
    local budget="${2:-${MAX_SESSION_COST:-1.00}}"

    local remaining
    remaining=$(get_remaining_budget "$current_cost")
    local percent_used
    percent_used=$(echo "scale=1; $current_cost / $budget * 100" | bc 2>/dev/null || echo "0")

    echo "$(format_cost "$current_cost") / $(format_cost "$budget") (${percent_used}% used)"
}

# Enforce budget and take action based on BUDGET_ACTION
# Returns: 0 = proceed, 1 = stop
# Usage: enforce_budget <current_cost> <additional_estimate> <context_msg>
enforce_budget() {
    local current_cost="${1:-0}"
    local additional_estimate="${2:-0}"
    local context_msg="${3:-operation}"
    local budget="${MAX_SESSION_COST:-1.00}"
    local action="${BUDGET_ACTION:-warn}"

    # Skip if budget enforcement is disabled
    if ! is_budget_enabled; then
        return 0
    fi

    # Calculate projected cost
    local projected_cost
    projected_cost=$(echo "scale=6; $current_cost + $additional_estimate" | bc)

    # Check if projected cost exceeds budget
    if ! check_budget "$projected_cost" "$budget"; then
        local msg="Budget limit exceeded for $context_msg: projected $(format_cost "$projected_cost") > $(format_cost "$budget")"

        case "$action" in
            stop)
                log_error "$msg"
                log_error "Stopping consultation (BUDGET_ACTION=stop)"
                return 1
                ;;
            warn|*)
                log_warn "$msg"
                log_warn "Continuing despite budget limit (BUDGET_ACTION=warn)"
                return 0
                ;;
        esac
    fi

    return 0
}

# Estimate cost for a specific phase
# Usage: estimate_phase_cost <phase> <num_consultants> <context_size>
estimate_phase_cost() {
    local phase="$1"
    local num_consultants="${2:-4}"
    local context_size="${3:-5000}"

    case "$phase" in
        round1|consultation)
            # Initial consultation: all consultants
            estimate_consultation_cost "$num_consultants" "$context_size"
            ;;
        synthesis)
            # Synthesis: single model, larger output
            estimate_query_cost "claude-3-sonnet" 2000 1500
            ;;
        *)
            echo "0"
            ;;
    esac
}