#!/bin/bash

# Status line with daily token usage tracking for Claude Code
input=$(cat)

# Extract basic info
cwd=$(echo "$input" | jq -r '.workspace.current_dir')
transcript_path=$(echo "$input" | jq -r '.transcript_path')
base_dir=$(basename "$cwd")

# Git info
git_info=""
if git -C "$cwd" rev-parse --git-dir > /dev/null 2>&1; then
  branch=$(git -C "$cwd" --no-optional-locks rev-parse --abbrev-ref HEAD 2>/dev/null)
  if ! git -C "$cwd" --no-optional-locks diff --quiet 2>/dev/null || \
     ! git -C "$cwd" --no-optional-locks diff --cached --quiet 2>/dev/null; then
    git_info=$(printf "\033[1;34mgit:(\033[0;31m%s\033[1;34m) \033[0;33m✗\033[0m " "$branch")
  else
    git_info=$(printf "\033[1;34mgit:(\033[0;31m%s\033[1;34m)\033[0m " "$branch")
  fi
fi

# Daily token tracking file
DAILY_TOKENS_FILE="$HOME/.claude/daily_tokens.json"
TODAY=$(date +%Y-%m-%d)

# Initialize or load daily tokens data
if [ ! -f "$DAILY_TOKENS_FILE" ]; then
  # Create initial structure
  echo "{\"date\":\"$TODAY\",\"sessions\":{}}" > "$DAILY_TOKENS_FILE"
fi

# Check if we need to reset for a new day
stored_date=$(jq -r '.date // ""' "$DAILY_TOKENS_FILE" 2>/dev/null)
if [ "$stored_date" != "$TODAY" ]; then
  # New day, reset the file
  echo "{\"date\":\"$TODAY\",\"sessions\":{}}" > "$DAILY_TOKENS_FILE"
fi

# Token info - parse JSONL format
token_info=""
if [ -f "$transcript_path" ]; then
  # Calculate tokens from current session's JSONL file - separate by type
  session_input_tokens=$(jq -s '
    [.[] | select(.type == "assistant") |
     .message.usage.input_tokens // 0
    ] | add // 0
  ' "$transcript_path" 2>/dev/null)

  session_cache_creation_tokens=$(jq -s '
    [.[] | select(.type == "assistant") |
     .message.usage.cache_creation_input_tokens // 0
    ] | add // 0
  ' "$transcript_path" 2>/dev/null)

  session_cache_read_tokens=$(jq -s '
    [.[] | select(.type == "assistant") |
     .message.usage.cache_read_input_tokens // 0
    ] | add // 0
  ' "$transcript_path" 2>/dev/null)

  session_output_tokens=$(jq -s '
    [.[] | select(.type == "assistant") |
     .message.usage.output_tokens // 0
    ] | add // 0
  ' "$transcript_path" 2>/dev/null)

  # Update session tokens in the daily file (overwrite, not add)
  if [ -n "$session_input_tokens" ] && [ -n "$session_output_tokens" ] && \
     [ "$session_input_tokens" != "null" ] && [ "$session_output_tokens" != "null" ]; then

    # Escape the transcript path for use as JSON key
    escaped_path=$(echo "$transcript_path" | sed 's/\\/\\\\/g; s/"/\\"/g')

    # Update the session data with all token types
    jq --arg path "$escaped_path" \
       --argjson inp "$session_input_tokens" \
       --argjson cr "$session_cache_creation_tokens" \
       --argjson rd "$session_cache_read_tokens" \
       --argjson outp "$session_output_tokens" \
       '.sessions[$path] = {"input": $inp, "cache_create": $cr, "cache_read": $rd, "output": $outp}' \
       "$DAILY_TOKENS_FILE" > "${DAILY_TOKENS_FILE}.tmp" && \
       mv "${DAILY_TOKENS_FILE}.tmp" "$DAILY_TOKENS_FILE"

    # Calculate total from all sessions - separate by type
    input_tokens=$(jq '[.sessions[].input] | add // 0' "$DAILY_TOKENS_FILE" 2>/dev/null)
    cache_create_tokens=$(jq '[.sessions[].cache_create] | add // 0' "$DAILY_TOKENS_FILE" 2>/dev/null)
    cache_read_tokens=$(jq '[.sessions[].cache_read] | add // 0' "$DAILY_TOKENS_FILE" 2>/dev/null)
    output_tokens=$(jq '[.sessions[].output] | add // 0' "$DAILY_TOKENS_FILE" 2>/dev/null)

    # Format and display if we have valid data
    if [ "$input_tokens" != "0" ] || [ "$cache_create_tokens" != "0" ] || [ "$cache_read_tokens" != "0" ] || [ "$output_tokens" != "0" ]; then
      total_tokens=$((input_tokens + cache_create_tokens + cache_read_tokens + output_tokens))

      # Function to format number with K (thousands) and M (millions)
      format_number() {
        local num=$1
        if [ "$num" -ge 1000000 ]; then
          # Convert to M with 1 decimal place
          local m_value=$(echo "scale=1; $num / 1000000" | bc | sed 's/\.0$//')
          echo "${m_value}M"
        elif [ "$num" -ge 1000 ]; then
          # Convert to K with 1 decimal place
          local k_value=$(echo "scale=1; $num / 1000" | bc | sed 's/\.0$//')
          echo "${k_value}K"
        else
          echo "$num"
        fi
      }

      # Format with K/M suffix for all token types
      input_fmt=$(format_number $input_tokens)
      cache_create_fmt=$(format_number $cache_create_tokens)
      cache_read_fmt=$(format_number $cache_read_tokens)
      output_fmt=$(format_number $output_tokens)
      total_fmt=$(format_number $total_tokens)

      token_info=$(printf "\033[0;35m[📊 今日用量 |\033[0m In:\033[0;36m%s\033[0m CW:\033[0;33m%s\033[0m CR:\033[0;32m%s\033[0m Out:\033[0;35m%s\033[0m Total:\033[1;32m%s\033[0;35m]\033[0m " \
        "$input_fmt" "$cache_create_fmt" "$cache_read_fmt" "$output_fmt" "$total_fmt")
    fi
  fi
fi

# Output final status line
printf "\033[1;32m➜\033[0m  \033[0;36m%s\033[0m %s\n%s" "$base_dir" "$git_info" "$token_info"