#!/usr/bin/env bash

# Session cost thresholds (USD) — on a $20 subscription these signal session weight,
# not real spend. Calibrate upward if you rarely see amber.
COST_AMBER=1.50
COST_RED=3.00

input=$(cat)

# --- Raw values from Claude Code ---
working_dir=$(echo "$input"       | jq -r '.workspace.current_dir // .cwd // empty')
model_name=$(echo "$input"        | jq -r '.model.display_name // empty')
context_used_pct=$(echo "$input"  | jq -r '.context_window.used_percentage // empty')
cost_usd=$(echo "$input"          | jq -r '.cost.total_cost_usd // empty')
input_tokens=$(echo "$input"      | jq -r '.context_window.total_input_tokens // 0')
output_tokens=$(echo "$input"     | jq -r '.context_window.total_output_tokens // 0')
vim_mode=$(echo "$input"          | jq -r '.vim.mode // empty')

dir_name=$(basename "$working_dir")

git_branch=""
if [ -n "$working_dir" ] && git -C "$working_dir" rev-parse --is-inside-work-tree 2>/dev/null | grep -q true; then
  git_branch=$(git -C "$working_dir" -c core.fsync=none symbolic-ref --short HEAD 2>/dev/null)
fi

# --- Colors ($'...' interprets \033 as actual ESC character) ---
COLOR_RESET=$'\033[0m'
COLOR_BOLD_BLUE=$'\033[1;34m'  # dir — identity, always stable
COLOR_BLUE=$'\033[0;34m'        # tokens — informational, no urgency
COLOR_CYAN=$'\033[0;36m'        # model — configuration, always stable
COLOR_GREEN=$'\033[0;32m'       # healthy / within limits
COLOR_AMBER=$'\033[0;33m'       # attention / approaching a limit
COLOR_RED=$'\033[0;31m'         # urgent / act now
COLOR_MAGENTA=$'\033[0;35m'     # vim mode — modal state indicator

# --- Branch: red on main/master (you are on trunk — danger), green on feature branches ---
branch_color="$COLOR_GREEN"
if [[ "$git_branch" == "main" || "$git_branch" == "master" ]]; then
  branch_color="$COLOR_RED"
fi

# --- Model tier: color and context thresholds both reflect cost per token.
#     Opus ~$15/MTok → red label, warn at 55/75%
#     Sonnet ~$3/MTok → amber label (mid-tier), warn at 65/85%
#     Haiku ~$0.80/MTok → green label, can fill further at 75/90%
if   [[ "$model_name" == *"Opus"* ]];  then model_color="$COLOR_RED";   CTX_AMBER=55; CTX_RED=75
elif [[ "$model_name" == *"Haiku"* ]]; then model_color="$COLOR_GREEN"; CTX_AMBER=75; CTX_RED=90
else                                        model_color="$COLOR_AMBER"; CTX_AMBER=65; CTX_RED=85
fi

# --- Token count: cumulative session total, formatted as Xk when >= 1000 ---
total_tokens=$(( input_tokens + output_tokens ))
if [ "$total_tokens" -ge 1000 ]; then
  tokens_display=$(awk "BEGIN { printf \"%.0fk\", $total_tokens / 1000 }")
else
  tokens_display="$total_tokens"
fi

# --- Context color: driven by model-aware thresholds above ---
context_color="$COLOR_GREEN"
context_pct=0
if [ -n "$context_used_pct" ]; then
  context_pct=$(printf '%.0f' "$context_used_pct")
  [ "$context_pct" -ge "$CTX_AMBER" ] && context_color="$COLOR_AMBER"
  [ "$context_pct" -ge "$CTX_RED" ]   && context_color="$COLOR_RED"
fi

# --- Cost: color and display format ---
cost_color="$COLOR_GREEN"
cost_display=""
if [ -n "$cost_usd" ] && [ "$cost_usd" != "0" ]; then
  cost_color=$(awk -v c="$cost_usd" -v amber="$COST_AMBER" -v red="$COST_RED" 'BEGIN {
    if (c >= red)        print "\033[0;31m"
    else if (c >= amber) print "\033[0;33m"
    else                 print "\033[0;32m"
  }')

  # Show cents (¢) below $1 for precision; dollars above
  cost_display=$(awk -v c="$cost_usd" 'BEGIN {
    if (c < 1.00) printf "%.0f\xC2\xA2", c * 100
    else          printf "$%.2f", c
  }')
fi

# --- Assemble: identity (where) → config (what) → health (context → cost → tokens) ---
parts=()
[ -n "$dir_name" ]         && parts+=("${COLOR_BOLD_BLUE}${dir_name}${COLOR_RESET}")
[ -n "$git_branch" ]       && parts+=("${branch_color}${git_branch}${COLOR_RESET}")
[ -n "$model_name" ]       && parts+=("${model_color}${model_name}${COLOR_RESET}")
[ -n "$context_used_pct" ] && parts+=("${context_color}context ${context_pct}%${COLOR_RESET}")
[ -n "$cost_display" ]     && parts+=("${cost_color}${cost_display}${COLOR_RESET}")
[ "$total_tokens" -gt 0 ]  && parts+=("${COLOR_BLUE}${tokens_display} tokens${COLOR_RESET}")
[ -n "$vim_mode" ]         && parts+=("${COLOR_MAGENTA}${vim_mode}${COLOR_RESET}")

# --- Join with ' · ' separator and print ---
output=""
for part in "${parts[@]}"; do
  [ -z "$output" ] && output="$part" || output="${output} · ${part}"
done

printf '%s' "$output"