#!/bin/bash

# Claude Code Configuration Script for API Router
# This script configures Claude Code to use your API Router instance

set -e

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

# Configuration
DEFAULT_BASE_URL="https://co.yes.vg"
CLAUDE_CONFIG_DIR="$HOME/.claude"
CLAUDE_SETTINGS_FILE="$CLAUDE_CONFIG_DIR/settings.json"

# Function to print colored output
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Function to check if jq is installed
check_jq() {
    if ! command -v jq &> /dev/null; then
        print_error "jq is required but not installed."
        print_info "Please install jq:"
        print_info "  macOS: brew install jq"
        print_info "  Ubuntu/Debian: sudo apt-get install jq"
        print_info "  CentOS/RHEL: sudo yum install jq"
        exit 1
    fi
}

# Function to backup existing settings
backup_settings() {
    if [ -f "$CLAUDE_SETTINGS_FILE" ]; then
        local backup_file="${CLAUDE_SETTINGS_FILE}.backup.$(date +%Y%m%d_%H%M%S)"
        cp "$CLAUDE_SETTINGS_FILE" "$backup_file"
        print_info "Backed up existing settings to: $backup_file"
    fi
}

# Function to create settings directory
create_settings_dir() {
    if [ ! -d "$CLAUDE_CONFIG_DIR" ]; then
        mkdir -p "$CLAUDE_CONFIG_DIR"
        print_info "Created Claude configuration directory: $CLAUDE_CONFIG_DIR"
    fi
}

# Function to validate API key format
validate_api_key() {
    local api_key="$1"
    if [[ ! "$api_key" =~ ^[A-Za-z0-9_-]+$ ]]; then
        print_error "Invalid API key format. API key should contain only alphanumeric characters, hyphens, and underscores."
        return 1
    fi
    return 0
}

api_root_url() {
    local base_url="${1%/}"
    case "$base_url" in
        */team/gemini)
            echo "${base_url%/team/gemini}"
            ;;
        */gemini/team)
            echo "${base_url%/gemini/team}"
            ;;
        */gemini)
            echo "${base_url%/gemini}"
            ;;
        */team)
            echo "${base_url%/team}"
            ;;
        *)
            echo "$base_url"
            ;;
    esac
}

is_team_url() {
    local base_url="${1%/}"
    [[ "$base_url" == */team || "$base_url" == */team/gemini || "$base_url" == */gemini/team ]]
}

# Function to test API connection
test_api_connection() {
    local base_url="$1"
    local api_key="$2"
    
    print_info "Testing API connection..."
    
    # Team vs user balance with legacy fallback
    local test_endpoint
    local balance_field
    local root_url
    root_url="$(api_root_url "$base_url")"
    if is_team_url "$base_url"; then
        test_endpoint="$root_url/api/v1/team/stats/spending"
        balance_field="daily_remaining"
    else
        test_endpoint="$root_url/api/v1/user/balance"
        balance_field="balance"
    fi
    
    # Test a simple request to the API
    local response
    response=$(curl --connect-timeout 5 --max-time 10 -s -w "%{http_code}" -o /tmp/claude_test_response \
        -X GET "$test_endpoint" \
        -H "Content-Type: application/json" \
        -H "X-API-Key: $api_key" \
        2>/dev/null || echo "000")
    
    if [ "$response" = "200" ]; then
        local balance
        balance=$(cat /tmp/claude_test_response | jq -r ".${balance_field}" 2>/dev/null || echo "unknown")
        if is_team_url "$base_url"; then
            print_success "API connection successful! Daily remaining: \$${balance}"
        else
            print_success "API connection successful! Current balance: \$${balance}"
        fi
        rm -f /tmp/claude_test_response
        return 0
    elif [ "$response" = "401" ]; then
        print_error "API key authentication failed. Please check your API key."
        rm -f /tmp/claude_test_response
        return 1
    elif [ "$response" = "000" ]; then
        print_error "Cannot connect to API server. Please check the URL and your internet connection."
        rm -f /tmp/claude_test_response
        return 1
    else
        print_error "API test failed with HTTP status: $response"
        rm -f /tmp/claude_test_response
        return 1
    fi
}

# Function to create or update Claude Code settings
create_settings() {
    local base_url="$1"
    local api_key="$2"

    local settings_json

    # Check if settings file exists and is valid JSON
    if [ -f "$CLAUDE_SETTINGS_FILE" ]; then
        # Try to read existing settings
        if jq . "$CLAUDE_SETTINGS_FILE" > /dev/null 2>&1; then
            print_info "Updating existing settings file (preserving other settings)..."

            # Merge: update only the YesCode endpoint and token values
            settings_json=$(jq --arg url "$base_url" --arg token "$api_key" '
                # Ensure env object exists
                .env = (.env // {}) |
                # Update only the API URL and token, and remove stale custom OAuth URL
                .env.ANTHROPIC_BASE_URL = $url |
                del(.env.CLAUDE_CODE_CUSTOM_OAUTH_URL) |
                .env.ANTHROPIC_AUTH_TOKEN = $token
            ' "$CLAUDE_SETTINGS_FILE")
        else
            print_warning "Existing settings file is invalid JSON, creating new one..."
            settings_json=$(cat <<EOF
{
  "env": {
    "ANTHROPIC_BASE_URL": "$base_url",
    "ANTHROPIC_AUTH_TOKEN": "$api_key",
    "DISABLE_TELEMETRY": 1,
    "DISABLE_ERROR_REPORTING": 1,
    "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": 1,
    "CLAUDE_BASH_MAINTAIN_PROJECT_WORKING_DIR": 1,
    "MAX_THINKING_TOKENS": 12000
  }
}
EOF
            )
        fi
    else
        # No existing file, create new one with defaults
        print_info "Creating new settings file with defaults..."
        settings_json=$(cat <<EOF
{
  "env": {
    "ANTHROPIC_BASE_URL": "$base_url",
    "ANTHROPIC_AUTH_TOKEN": "$api_key",
    "DISABLE_TELEMETRY": 1,
    "DISABLE_ERROR_REPORTING": 1,
    "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": 1,
    "CLAUDE_BASH_MAINTAIN_PROJECT_WORKING_DIR": 1,
    "MAX_THINKING_TOKENS": 12000
  }
}
EOF
        )
    fi

    # Validate JSON
    if ! echo "$settings_json" | jq . > /dev/null 2>&1; then
        print_error "Generated settings JSON is invalid"
        return 1
    fi

    # Write settings file
    echo "$settings_json" > "$CLAUDE_SETTINGS_FILE"
    print_success "Claude Code settings written to: $CLAUDE_SETTINGS_FILE"
}

remove_custom_oauth_env_from_file() {
    local shell_config="$1"
    if [ -f "$shell_config" ] && grep -q "CLAUDE_CODE_CUSTOM_OAUTH_URL" "$shell_config"; then
        if [[ "$OSTYPE" == "darwin"* ]]; then
            sed -i '' '/CLAUDE_CODE_CUSTOM_OAUTH_URL/d' "$shell_config"
        else
            sed -i '/CLAUDE_CODE_CUSTOM_OAUTH_URL/d' "$shell_config"
        fi
        print_info "Removed CLAUDE_CODE_CUSTOM_OAUTH_URL from $shell_config"
    fi
}

# Function to set environment variables in shell profile
set_environment_variables() {
    local base_url="$1"
    local api_key="$2"

    # Export for current session
    export ANTHROPIC_BASE_URL="$base_url"
    export ANTHROPIC_AUTH_TOKEN="$api_key"
    unset CLAUDE_CODE_CUSTOM_OAUTH_URL

    # Detect shell and add to appropriate config file
    local shell_config=""
    local shell_name=""

    # First check $SHELL to determine user's default shell
    if [ -n "$SHELL" ]; then
        shell_name=$(basename "$SHELL")
        case "$shell_name" in
            bash)
                shell_config="$HOME/.bashrc"
                [ -f "$HOME/.bash_profile" ] && shell_config="$HOME/.bash_profile"
                ;;
            zsh)
                shell_config="$HOME/.zshrc"
                ;;
            fish)
                shell_config="$HOME/.config/fish/config.fish"
                ;;
            *)
                shell_config="$HOME/.profile"
                ;;
        esac
    # Fallback to checking version variables if $SHELL is not set
    elif [ -n "$BASH_VERSION" ]; then
        shell_config="$HOME/.bashrc"
        [ -f "$HOME/.bash_profile" ] && shell_config="$HOME/.bash_profile"
    elif [ -n "$ZSH_VERSION" ]; then
        shell_config="$HOME/.zshrc"
    elif [ -n "$FISH_VERSION" ]; then
        shell_config="$HOME/.config/fish/config.fish"
    else
        shell_config="$HOME/.profile"
    fi

    print_info "Detected shell: ${shell_name:-$(basename $SHELL 2>/dev/null || echo 'unknown')}"
    print_info "Using config file: $shell_config"

    remove_custom_oauth_env_from_file "$shell_config"

    # Handle Fish shell differently (uses 'set -x' instead of 'export')
    if [ "$shell_name" = "fish" ] || [[ "$shell_config" == *"fish"* ]]; then
        # Fish shell syntax
        mkdir -p "$(dirname "$shell_config")"

        # Update or add ANTHROPIC_BASE_URL
        if [ -f "$shell_config" ] && grep -q "set -x ANTHROPIC_BASE_URL" "$shell_config"; then
            if [[ "$OSTYPE" == "darwin"* ]]; then
                sed -i '' "s|set -x ANTHROPIC_BASE_URL.*|set -x ANTHROPIC_BASE_URL \"$base_url\"|" "$shell_config"
            else
                sed -i "s|set -x ANTHROPIC_BASE_URL.*|set -x ANTHROPIC_BASE_URL \"$base_url\"|" "$shell_config"
            fi
            print_info "Updated ANTHROPIC_BASE_URL in $shell_config"
        else
            echo "" >> "$shell_config"
            echo "# Claude Code API Router configuration" >> "$shell_config"
            echo "set -x ANTHROPIC_BASE_URL \"$base_url\"" >> "$shell_config"
            print_info "Added ANTHROPIC_BASE_URL to $shell_config"
        fi

        # Update or add ANTHROPIC_AUTH_TOKEN
        if [ -f "$shell_config" ] && grep -q "set -x ANTHROPIC_AUTH_TOKEN" "$shell_config"; then
            if [[ "$OSTYPE" == "darwin"* ]]; then
                sed -i '' "s|set -x ANTHROPIC_AUTH_TOKEN.*|set -x ANTHROPIC_AUTH_TOKEN \"$api_key\"|" "$shell_config"
            else
                sed -i "s|set -x ANTHROPIC_AUTH_TOKEN.*|set -x ANTHROPIC_AUTH_TOKEN \"$api_key\"|" "$shell_config"
            fi
            print_info "Updated ANTHROPIC_AUTH_TOKEN in $shell_config"
        else
            echo "set -x ANTHROPIC_AUTH_TOKEN \"$api_key\"" >> "$shell_config"
            print_info "Added ANTHROPIC_AUTH_TOKEN to $shell_config"
        fi
    else
        # Bash/Zsh/sh syntax
        # Update or add ANTHROPIC_BASE_URL
        if [ -f "$shell_config" ] && grep -q "export ANTHROPIC_BASE_URL=" "$shell_config"; then
            if [[ "$OSTYPE" == "darwin"* ]]; then
                sed -i '' "s|export ANTHROPIC_BASE_URL=.*|export ANTHROPIC_BASE_URL=\"$base_url\"|" "$shell_config"
            else
                sed -i "s|export ANTHROPIC_BASE_URL=.*|export ANTHROPIC_BASE_URL=\"$base_url\"|" "$shell_config"
            fi
            print_info "Updated ANTHROPIC_BASE_URL in $shell_config"
        else
            echo "" >> "$shell_config"
            echo "# Claude Code API Router configuration" >> "$shell_config"
            echo "export ANTHROPIC_BASE_URL=\"$base_url\"" >> "$shell_config"
            print_info "Added ANTHROPIC_BASE_URL to $shell_config"
        fi

        # Update or add ANTHROPIC_AUTH_TOKEN
        if [ -f "$shell_config" ] && grep -q "export ANTHROPIC_AUTH_TOKEN=" "$shell_config"; then
            if [[ "$OSTYPE" == "darwin"* ]]; then
                sed -i '' "s|export ANTHROPIC_AUTH_TOKEN=.*|export ANTHROPIC_AUTH_TOKEN=\"$api_key\"|" "$shell_config"
            else
                sed -i "s|export ANTHROPIC_AUTH_TOKEN=.*|export ANTHROPIC_AUTH_TOKEN=\"$api_key\"|" "$shell_config"
            fi
            print_info "Updated ANTHROPIC_AUTH_TOKEN in $shell_config"
        else
            echo "export ANTHROPIC_AUTH_TOKEN=\"$api_key\"" >> "$shell_config"
            print_info "Added ANTHROPIC_AUTH_TOKEN to $shell_config"
        fi
    fi

    return 0
}

# Function to display current settings
display_settings() {
    if [ -f "$CLAUDE_SETTINGS_FILE" ]; then
        print_info "Current Claude Code settings:"
        echo "----------------------------------------"
        cat "$CLAUDE_SETTINGS_FILE" | jq .
        echo "----------------------------------------"
    else
        print_info "No existing Claude Code settings found."
    fi

    echo ""
    print_info "Current environment variables:"
    echo "----------------------------------------"

    if [ -n "$ANTHROPIC_BASE_URL" ]; then
        print_info "ANTHROPIC_BASE_URL: $ANTHROPIC_BASE_URL"
    else
        print_info "ANTHROPIC_BASE_URL: (not set)"
    fi

    if [ -n "$ANTHROPIC_AUTH_TOKEN" ]; then
        local masked_token
        if [ ${#ANTHROPIC_AUTH_TOKEN} -gt 12 ]; then
            masked_token="${ANTHROPIC_AUTH_TOKEN:0:8}...${ANTHROPIC_AUTH_TOKEN: -4}"
        else
            masked_token="${ANTHROPIC_AUTH_TOKEN:0:4}..."
        fi
        print_info "ANTHROPIC_AUTH_TOKEN: $masked_token"
    else
        print_info "ANTHROPIC_AUTH_TOKEN: (not set)"
    fi
    echo "----------------------------------------"
}

# Main function
main() {
    print_info "Claude Code Configuration Script for API Router"
    echo "======================================================="
    echo
    
    # Check dependencies
    check_jq
    
    # Parse command line arguments
    local base_url=""
    local api_key=""
    local test_only=false
    local show_settings=false
    
    while [[ $# -gt 0 ]]; do
        case $1 in
            -u|--url)
                base_url="$2"
                shift 2
                ;;
            -k|--key)
                api_key="$2"
                shift 2
                ;;
            -t|--test)
                test_only=true
                shift
                ;;
            -s|--show)
                show_settings=true
                shift
                ;;
            -h|--help)
                cat <<EOF
Usage: $0 [OPTIONS]

Options:
  -u, --url URL     Set the API Router base URL (default: $DEFAULT_BASE_URL)
  -k, --key KEY     Set the API key
  -t, --test        Test API connection only (requires -u and -k)
  -s, --show        Show current settings and exit
  -h, --help        Show this help message

Examples:
  $0 --url https://your-domain.tld --key your-api-key-here
  $0 --test --url https://your-domain.tld --key your-api-key-here
  $0 --show

Interactive mode (no arguments):
  $0
EOF
                exit 0
                ;;
            *)
                print_error "Unknown option: $1"
                print_info "Use --help for usage information"
                exit 1
                ;;
        esac
    done
    
    # Show settings and exit if requested
    if [ "$show_settings" = true ]; then
        display_settings
        exit 0
    fi
    
    # Interactive mode if no arguments provided
    if [ -z "$base_url" ] && [ -z "$api_key" ]; then
        print_info "Interactive setup mode"
        echo
        
        # Get base URL
        read -p "Enter API Router URL [$DEFAULT_BASE_URL]: " base_url
        if [ -z "$base_url" ]; then
            base_url="$DEFAULT_BASE_URL"
        fi
        
        # Get API key
        while [ -z "$api_key" ]; do
            read -p "Enter your API key: " api_key
            if [ -z "$api_key" ]; then
                print_warning "API key is required"
            elif ! validate_api_key "$api_key"; then
                api_key=""
            fi
        done
    fi

    if [ -z "$base_url" ]; then
        base_url="$DEFAULT_BASE_URL"
    fi
    
    # Validate inputs
    if [ -z "$base_url" ] || [ -z "$api_key" ]; then
        print_error "Both URL and API key are required"
        print_info "Use --help for usage information"
        exit 1
    fi
    
    # Validate API key
    if ! validate_api_key "$api_key"; then
        exit 1
    fi
    
    # Remove trailing slash from URL
    base_url="${base_url%/}"
    
    print_info "Configuration:"
    print_info "  Base URL: $base_url"
    print_info "  API Key: ${api_key:0:8}...${api_key: -4}"
    echo
    
    # Test API connection
    if ! test_api_connection "$base_url" "$api_key"; then
        print_warning "API check failed; continuing without connection verification."
    fi
    
    # Exit if test only
    if [ "$test_only" = true ]; then
        print_info "API check completed (diagnostic only)."
        exit 0
    fi
    
    # Create settings directory
    create_settings_dir
    
    # Backup existing settings
    backup_settings
    
    # Create new settings
    if create_settings "$base_url" "$api_key"; then
        echo

        # Set environment variables in shell profile
        print_info "Setting environment variables..."
        if set_environment_variables "$base_url" "$api_key"; then
            print_success "Environment variables set successfully"
        else
            print_warning "Failed to set environment variables automatically"
            print_info "Please set manually:"
            print_info "  export ANTHROPIC_BASE_URL=\"$base_url\""
            print_info "  export ANTHROPIC_AUTH_TOKEN=\"$api_key\""
        fi

        echo
        print_success "Claude Code has been configured successfully!"
        print_info "You can now use Claude Code with your API router."
        print_info ""
        print_info "To apply the environment variables in your current session, run:"

        # Provide correct command based on detected shell
        local current_shell=$(basename "$SHELL" 2>/dev/null || echo "bash")
        if [ "$current_shell" = "fish" ]; then
            print_info "  set -x ANTHROPIC_BASE_URL \"$base_url\""
            print_info "  set -x ANTHROPIC_AUTH_TOKEN \"$api_key\""
        else
            print_info "  export ANTHROPIC_BASE_URL=\"$base_url\""
            print_info "  export ANTHROPIC_AUTH_TOKEN=\"$api_key\""
        fi
        print_info "Or restart your terminal."
        print_info ""
        print_info "To verify the setup, run:"
        print_info "  yescode --version"
        print_info ""
        print_info "Configuration file location: $CLAUDE_SETTINGS_FILE"
        print_info "Environment variables have been set for ANTHROPIC_BASE_URL and ANTHROPIC_AUTH_TOKEN"

        if [ -f "$CLAUDE_SETTINGS_FILE" ]; then
            echo
            print_info "Current settings:"
            cat "$CLAUDE_SETTINGS_FILE" | jq .
        fi
    else
        print_error "Failed to create Claude Code settings"
        exit 1
    fi
}

# Run main function
main "$@"