#!/bin/bash

# Read JSON input from stdin
input=$(cat)

# Extract data from JSON
current_dir=$(echo "$input" | jq -r '.workspace.current_dir // empty')
model=$(echo "$input" | jq -r '.model.display_name // "?"')
total_input=$(echo "$input" | jq -r '.context_window.total_input_tokens // 0')
total_output=$(echo "$input" | jq -r '.context_window.total_output_tokens // 0')
context_size=$(echo "$input" | jq -r '.context_window.context_window_size // 200000')
used_percent=$(echo "$input" | jq -r '.context_window.used_percentage // 0' | cut -d'.' -f1)
cost_usd=$(echo "$input" | jq -r '.cost.total_cost_usd // 0')

# Calculate total tokens in K format
total_tokens=$(( (total_input + total_output) / 1000 ))

# Format cost (show cents if < $1, otherwise dollars)
# Use LC_NUMERIC=C to ensure decimal point handling
cost_cents=$(LC_NUMERIC=C awk "BEGIN {printf \"%.0f\", $cost_usd * 100}")
if [ "$cost_cents" -lt 100 ]; then
    cost_formatted="${cost_cents}¢"
else
    cost_dollars=$(LC_NUMERIC=C awk "BEGIN {printf \"%.2f\", $cost_usd}")
    cost_formatted="\$${cost_dollars}"
fi

# Get user and hostname
user=$(whoami)
host=$(hostname -s)

# Get git branch if in a git repository
git_branch=""
if [ -n "$current_dir" ] && [ -d "$current_dir/.git" ] || git -C "$current_dir" rev-parse --git-dir > /dev/null 2>&1; then
    branch=$(git -C "$current_dir" branch --show-current 2>/dev/null)
    [ -n "$branch" ] && git_branch=" \033[33m($branch)\033[00m"
fi

# Check for debian_chroot
chroot_prefix=""
if [ -r /etc/debian_chroot ]; then
    debian_chroot=$(cat /etc/debian_chroot)
    chroot_prefix="(${debian_chroot})"
fi

# Format context limit in K
context_limit_k=$(( context_size / 1000 ))

# Output status line
# Format: user@host:dir (branch) | Model | Tokens/Limit | Cost
printf "${chroot_prefix}\033[32m${user}@${host}\033[00m:\033[34m${current_dir}\033[00m${git_branch}"
printf " \033[90m|\033[00m \033[36m${model}\033[00m"
printf " \033[90m|\033[00m \033[35m${total_tokens}K/${context_limit_k}K (${used_percent}%%)\033[00m"
printf " \033[90m|\033[00m \033[33m${cost_formatted}\033[00m"