#!/bin/bash
# gemini-wrapper - Token-aware wrapper for gemini command
# Location: ~/.claude/scripts/gemini-wrapper
#
# This wrapper automatically manages --all-files flag based on project token count
# and provides intelligent approval mode defaults
#
# Usage: gemini-wrapper [all gemini options]
#
# Approval Mode Options:
#   --approval-mode default    : Prompt for approval on each tool call (default)
#   --approval-mode auto_edit  : Auto-approve edit tools, prompt for others
#   --approval-mode yolo       : Auto-approve all tool calls
#
# Note: Executes in current working directory

set -e

# Function to show help
show_help() {
    echo "gemini-wrapper - Token-aware wrapper for gemini command"
    echo ""
    echo "Usage: gemini-wrapper [options] [gemini options]"
    echo ""
    echo "Options:"
    echo "  --approval-mode <mode>  Sets the approval mode for tool calls"
    echo "                          Available modes:"
    echo "                            default    : Prompt for approval on each tool call (default)"
    echo "                            auto_edit  : Auto-approve edit tools, prompt for others"
    echo "                            yolo       : Auto-approve all tool calls"
    echo "  --help                  Show this help message"
    echo ""
    echo "Features:"
    echo "  - Automatically manages --all-files flag based on project token count"
    echo "  - Intelligent approval mode detection based on task type"
    echo "  - Token limit: $DEFAULT_TOKEN_LIMIT (set GEMINI_TOKEN_LIMIT to override)"
    echo ""
    echo "Examples:"
    echo "  gemini-wrapper -p \"Analyze the codebase structure\""
    echo "  gemini-wrapper --approval-mode yolo -p \"Implement user authentication\""
    echo "  gemini-wrapper --approval-mode auto_edit -p \"Fix all linting errors\""
    echo ""
}

# Configuration
DEFAULT_TOKEN_LIMIT=2000000
TOKEN_LIMIT=${GEMINI_TOKEN_LIMIT:-$DEFAULT_TOKEN_LIMIT}

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Function to count tokens (approximate: chars/4) - optimized version
count_tokens() {
    local total_chars=0
    local file_count=0

    # Use single find with bulk wc for better performance
    # Common source file extensions
    local extensions="py js ts tsx jsx java cpp c h rs go md txt json yaml yml xml html css scss sass php rb sh bash"

    # Build find command with extension patterns
    local find_cmd="find . -type f \("
    local first=true
    for ext in $extensions; do
        if [[ "$first" == true ]]; then
            find_cmd+=" -name \"*.$ext\""
            first=false
        else
            find_cmd+=" -o -name \"*.$ext\""
        fi
    done
    find_cmd+=" \)"

    # Exclude common build/cache directories
    find_cmd+=" -not -path \"*/node_modules/*\""
    find_cmd+=" -not -path \"*/.git/*\""
    find_cmd+=" -not -path \"*/dist/*\""
    find_cmd+=" -not -path \"*/build/*\""
    find_cmd+=" -not -path \"*/.next/*\""
    find_cmd+=" -not -path \"*/.nuxt/*\""
    find_cmd+=" -not -path \"*/target/*\""
    find_cmd+=" -not -path \"*/vendor/*\""
    find_cmd+=" -not -path \"*/__pycache__/*\""
    find_cmd+=" -not -path \"*/.cache/*\""
    find_cmd+=" 2>/dev/null"

    # Use efficient bulk processing with wc
    if command -v wc >/dev/null 2>&1; then
        # Try bulk wc first - much faster for many files
        local wc_output
        wc_output=$(eval "$find_cmd" | xargs wc -c 2>/dev/null | tail -n 1)

        # Parse the total line (last line of wc output when processing multiple files)
        if [[ -n "$wc_output" && "$wc_output" =~ ^[[:space:]]*([0-9]+)[[:space:]]+total[[:space:]]*$ ]]; then
            total_chars="${BASH_REMATCH[1]}"
            file_count=$(eval "$find_cmd" | wc -l 2>/dev/null || echo 0)
        else
            # Fallback: single file processing
            while IFS= read -r file; do
                if [[ -f "$file" && -r "$file" ]]; then
                    local chars=$(wc -c < "$file" 2>/dev/null || echo 0)
                    total_chars=$((total_chars + chars))
                    file_count=$((file_count + 1))
                fi
            done < <(eval "$find_cmd")
        fi
    else
        # No wc available - fallback method
        while IFS= read -r file; do
            if [[ -f "$file" && -r "$file" ]]; then
                local chars=$(stat -c%s "$file" 2>/dev/null || echo 0)
                total_chars=$((total_chars + chars))
                file_count=$((file_count + 1))
            fi
        done < <(eval "$find_cmd")
    fi

    local estimated_tokens=$((total_chars / 4))
    echo "$estimated_tokens $file_count"
}

# Function to validate approval mode
validate_approval_mode() {
    local mode="$1"
    case "$mode" in
        "default"|"auto_edit"|"yolo")
            return 0
            ;;
        *)
            echo -e "${RED}❌ Invalid approval mode: $mode${NC}" >&2
            echo -e "${YELLOW}Valid modes: default, auto_edit, yolo${NC}" >&2
            return 1
            ;;
    esac
}

# Parse arguments to check for flags
has_all_files=false
has_approval_mode=false
approval_mode_value=""
args=()
i=0

# Parse arguments with proper handling of --approval-mode value
args=("$@")  # Start with all arguments
parsed_args=()
skip_next=false

for ((i=0; i<${#args[@]}; i++)); do
    if [[ "$skip_next" == true ]]; then
        skip_next=false
        continue
    fi

    arg="${args[i]}"
    case "$arg" in
        "--help"|"-h")
            show_help
            exit 0
            ;;
        "--all-files")
            has_all_files=true
            parsed_args+=("$arg")
            ;;
        "--approval-mode")
            has_approval_mode=true
            # Get the next argument as the mode value
            if [[ $((i+1)) -lt ${#args[@]} ]]; then
                approval_mode_value="${args[$((i+1))]}"
                if validate_approval_mode "$approval_mode_value"; then
                    parsed_args+=("$arg" "$approval_mode_value")
                    skip_next=true
                else
                    exit 1
                fi
            else
                echo -e "${RED}❌ --approval-mode requires a value${NC}" >&2
                echo -e "${YELLOW}Valid modes: default, auto_edit, yolo${NC}" >&2
                exit 1
            fi
            ;;
        --approval-mode=*)
            has_approval_mode=true
            approval_mode_value="${arg#*=}"
            if validate_approval_mode "$approval_mode_value"; then
                parsed_args+=("$arg")
            else
                exit 1
            fi
            ;;
        *)
            parsed_args+=("$arg")
            ;;
    esac
done

# Replace args with parsed_args
args=("${parsed_args[@]}")

# Analyze current working directory
echo -e "${GREEN}📁 Analyzing current directory: $(pwd)${NC}" >&2

# Count tokens (in the target directory if -c was used)
echo -e "${YELLOW}🔍 Analyzing project size...${NC}" >&2
read -r token_count file_count <<< "$(count_tokens)"

echo -e "${YELLOW}📊 Project stats: ~${token_count} tokens across ${file_count} files${NC}" >&2

# Decision logic for --all-files flag
if [[ $token_count -lt $TOKEN_LIMIT ]]; then
    if [[ "$has_all_files" == false ]]; then
        echo -e "${GREEN}✅ Small project (${token_count} < ${TOKEN_LIMIT} tokens): Adding --all-files${NC}" >&2
        args=("--all-files" "${args[@]}")
    else
        echo -e "${GREEN}✅ Small project (${token_count} < ${TOKEN_LIMIT} tokens): Keeping --all-files${NC}" >&2
    fi
else
    if [[ "$has_all_files" == true ]]; then
        echo -e "${RED}⚠️  Large project (${token_count} >= ${TOKEN_LIMIT} tokens): Removing --all-files to avoid token limits${NC}" >&2
        echo -e "${YELLOW}💡 Consider using specific @{patterns} for targeted analysis${NC}" >&2
        # Remove --all-files from args
        new_args=()
        for arg in "${args[@]}"; do
            if [[ "$arg" != "--all-files" ]]; then
                new_args+=("$arg")
            fi
        done
        args=("${new_args[@]}")
    else
        echo -e "${RED}⚠️  Large project (${token_count} >= ${TOKEN_LIMIT} tokens): Avoiding --all-files${NC}" >&2
        echo -e "${YELLOW}💡 Consider using specific @{patterns} for targeted analysis${NC}" >&2
    fi
fi

# Auto-add approval-mode if not specified
if [[ "$has_approval_mode" == false ]]; then
    # Intelligent approval mode detection based on prompt content
    prompt_text="${args[*]}"

    # Analysis/Research tasks - use default (prompt for each tool)
    if [[ "$prompt_text" =~ (analyze|analysis|review|understand|inspect|examine|research|study|explore|investigate) ]]; then
        echo -e "${GREEN}📋 Analysis task detected: Adding --approval-mode default${NC}" >&2
        args=("--approval-mode" "default" "${args[@]}")

    # Development/Edit tasks - use auto_edit (auto-approve edits, prompt for others)
    elif [[ "$prompt_text" =~ (implement|create|build|develop|code|write|edit|modify|update|fix|refactor|generate) ]]; then
        echo -e "${GREEN}🔧 Development task detected: Adding --approval-mode auto_edit${NC}" >&2
        args=("--approval-mode" "auto_edit" "${args[@]}")

    # Automation/Batch tasks - use yolo (auto-approve all)
    elif [[ "$prompt_text" =~ (automate|batch|mass|bulk|all|execute|run|deploy|install|setup) ]]; then
        echo -e "${YELLOW}⚡ Automation task detected: Adding --approval-mode yolo${NC}" >&2
        args=("--approval-mode" "yolo" "${args[@]}")

    # Default fallback - use default mode for safety
    else
        echo -e "${YELLOW}🔍 General task detected: Adding --approval-mode default${NC}" >&2
        args=("--approval-mode" "default" "${args[@]}")
    fi

    # Show approval mode explanation
    case "${args[1]}" in
        "default")
            echo -e "${YELLOW}   → Will prompt for approval on each tool call${NC}" >&2
            ;;
        "auto_edit")
            echo -e "${YELLOW}   → Will auto-approve edit tools, prompt for others${NC}" >&2
            ;;
        "yolo")
            echo -e "${YELLOW}   → Will auto-approve all tool calls${NC}" >&2
            ;;
    esac
fi

# Show final command (for transparency)
echo -e "${YELLOW}🚀 Executing: gemini ${args[*]}${NC}" >&2

# Execute gemini with adjusted arguments (we're already in the right directory)
gemini "${args[@]}"