#!/bin/bash
# qwen-token-refresh.sh - Auto-refresh Qwen OAuth token
# Usage:
#   ./qwen-token-refresh.sh           # Check and refresh if expired
#   ./qwen-token-refresh.sh --daemon  # Run as background daemon
#   ./qwen-token-refresh.sh --status  # Check token status
#
# This enables auto token refresh for ALL platforms:
#   OpenClaw, NanoBot, PicoClaw, NanoClaw
#
# For ZeroClaw, this is NOT needed (native support).

set -e

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

CREDS_FILE="$HOME/.qwen/oauth_creds.json"
ENV_FILE="$HOME/.qwen/.env"
REFRESH_URL="https://chat.qwen.ai/api/v1/oauth2/token"
API_BASE_URL="https://dashscope.aliyuncs.com/compatible-mode/v1"

# Token refresh buffer (refresh 5 minutes before expiry)
REFRESH_BUFFER_SECONDS=300

check_dependencies() {
    if ! command -v jq &> /dev/null; then
        echo -e "${RED}Error: jq is required${NC}"
        echo "Install with: sudo apt install jq"
        exit 1
    fi

    if ! command -v curl &> /dev/null; then
        echo -e "${RED}Error: curl is required${NC}"
        exit 1
    fi
}

get_token_status() {
    if [ ! -f "$CREDS_FILE" ]; then
        echo "NOT_FOUND"
        return
    fi

    local expiry_date=$(cat "$CREDS_FILE" | jq -r '.expiry_date // 0')
    local current_time=$(date +%s)000

    if [ "$expiry_date" -eq 0 ]; then
        echo "UNKNOWN"
        return
    fi

    local expiry_seconds=$((expiry_date / 1000))
    local current_seconds=$((current_time / 1000))
    local remaining=$((expiry_seconds - current_seconds))

    if [ "$remaining" -lt 0 ]; then
        echo "EXPIRED"
    elif [ "$remaining" -lt "$REFRESH_BUFFER_SECONDS" ]; then
        echo "EXPIRING_SOON"
    else
        echo "VALID"
    fi
}

format_time_remaining() {
    local expiry_date=$1
    local expiry_seconds=$((expiry_date / 1000))
    local current_seconds=$(date +%s)
    local remaining=$((expiry_seconds - current_seconds))

    if [ "$remaining" -lt 0 ]; then
        echo "expired"
    elif [ "$remaining" -lt 60 ]; then
        echo "${remaining}s"
    elif [ "$remaining" -lt 3600 ]; then
        echo "$((remaining / 60))m"
    else
        echo "$((remaining / 3600))h $(((remaining % 3600) / 60))m"
    fi
}

refresh_token() {
    local refresh_token=$(cat "$CREDS_FILE" | jq -r '.refresh_token // empty')

    if [ -z "$refresh_token" ] || [ "$refresh_token" = "null" ]; then
        echo -e "${RED}Error: No refresh_token found in credentials${NC}"
        echo "Please re-authenticate: qwen --auth-type qwen-oauth -p 'test'"
        return 1
    fi

    echo -e "${BLUE}Refreshing token...${NC}"

    local response=$(curl -s -X POST "$REFRESH_URL" \
        -H "Content-Type: application/x-www-form-urlencoded" \
        -d "grant_type=refresh_token" \
        -d "refresh_token=$refresh_token" \
        2>/dev/null)

    if [ -z "$response" ]; then
        echo -e "${RED}Error: Empty response from refresh endpoint${NC}"
        return 1
    fi

    local new_access_token=$(echo "$response" | jq -r '.access_token // empty')
    local new_refresh_token=$(echo "$response" | jq -r '.refresh_token // empty')
    local expires_in=$(echo "$response" | jq -r '.expires_in // 3600')

    if [ -z "$new_access_token" ] || [ "$new_access_token" = "null" ]; then
        echo -e "${RED}Error: Failed to get new access_token${NC}"
        echo "Response: $response"
        echo ""
        echo "Please re-authenticate: qwen --auth-type qwen-oauth -p 'test'"
        return 1
    fi

    # Calculate new expiry date (current time + expires_in, in milliseconds)
    local current_ms=$(date +%s)000
    local expires_in_ms=$((expires_in * 1000))
    local new_expiry=$((current_ms + expires_in_ms))

    # Update credentials file
    local resource_url=$(cat "$CREDS_FILE" | jq -r '.resource_url // "portal.qwen.ai"')

    cat > "$CREDS_FILE" << EOF
{
  "access_token": "$new_access_token",
  "refresh_token": "${new_refresh_token:-$refresh_token}",
  "token_type": "Bearer",
  "resource_url": "$resource_url",
  "expiry_date": $new_expiry
}
EOF

    chmod 600 "$CREDS_FILE"

    # Update .env file for other platforms
    mkdir -p "$HOME/.qwen"
    cat > "$ENV_FILE" << EOF
# Qwen OAuth Configuration (auto-refreshed: $(date))
export OPENAI_API_KEY="$new_access_token"
export OPENAI_BASE_URL="$API_BASE_URL"
export OPENAI_MODEL="coder-model"
EOF

    chmod 600 "$ENV_FILE"

    echo -e "${GREEN}✅ Token refreshed successfully!${NC}"
    echo "   New token expires in: $(format_time_remaining $new_expiry)"
    echo ""
    echo "Environment updated. Run to apply:"
    echo "  source ~/.qwen/.env"

    return 0
}

show_status() {
    echo "🦞 Qwen OAuth Token Status"
    echo "==========================="
    echo ""

    if [ ! -f "$CREDS_FILE" ]; then
        echo -e "${RED}❌ Credentials file not found${NC}"
        echo "   Expected: $CREDS_FILE"
        echo ""
        echo "To authenticate, run:"
        echo "  qwen --auth-type qwen-oauth -p 'test'"
        exit 1
    fi

    local status=$(get_token_status)
    local expiry_date=$(cat "$CREDS_FILE" | jq -r '.expiry_date // 0')
    local access_token=$(cat "$CREDS_FILE" | jq -r '.access_token // ""')

    echo "Credentials: $CREDS_FILE"
    echo "Token: ${access_token:0:20}...${access_token: -10}"
    echo ""

    case "$status" in
        VALID)
            echo -e "Status: ${GREEN}✅ VALID${NC}"
            echo "Expires in: $(format_time_remaining $expiry_date)"
            ;;
        EXPIRING_SOON)
            echo -e "Status: ${YELLOW}⚠️  EXPIRING SOON${NC}"
            echo "Expires in: $(format_time_remaining $expiry_date)"
            echo ""
            echo "Run refresh: $0"
            ;;
        EXPIRED)
            echo -e "Status: ${RED}❌ EXPIRED${NC}"
            echo ""
            echo "Run refresh: $0"
            ;;
        UNKNOWN)
            echo -e "Status: ${YELLOW}⚠️  UNKNOWN${NC}"
            echo "Expiry date not found in credentials"
            ;;
    esac

    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "FREE tier: 2,000 requests/day, 60 req/min"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
}

run_daemon() {
    echo "🦞 Qwen Token Refresh Daemon"
    echo "============================="
    echo ""
    echo "Starting background refresh service..."
    echo "Check interval: 5 minutes"
    echo "Refresh buffer: ${REFRESH_BUFFER_SECONDS}s before expiry"
    echo ""
    echo "Press Ctrl+C to stop"
    echo ""

    while true; do
        local status=$(get_token_status)

        case "$status" in
            EXPIRED|EXPIRING_SOON)
                echo "$(date '+%Y-%m-%d %H:%M:%S') - Token needs refresh ($status)"
                if refresh_token > /dev/null 2>&1; then
                    echo "$(date '+%Y-%m-%d %H:%M:%S') - ✅ Token refreshed"
                else
                    echo "$(date '+%Y-%m-%d %H:%M:%S') - ❌ Refresh failed"
                fi
                ;;
            VALID)
                local expiry=$(cat "$CREDS_FILE" | jq -r '.expiry_date // 0')
                echo "$(date '+%Y-%m-%d %H:%M:%S') - Token valid ($(format_time_remaining $expiry) remaining)"
                ;;
            *)
                echo "$(date '+%Y-%m-%d %H:%M:%S') - Token status: $status"
                ;;
        esac

        sleep 300  # Check every 5 minutes
    done
}

install_systemd_service() {
    local service_dir="$HOME/.config/systemd/user"
    local service_file="$service_dir/qwen-token-refresh.service"

    mkdir -p "$service_dir"

    cat > "$service_file" << 'EOF'
[Unit]
Description=Qwen OAuth Token Auto-Refresh
After=network.target

[Service]
Type=simple
ExecStart=%h/.local/bin/qwen-token-refresh.sh --daemon
Restart=always
RestartSec=60

[Install]
WantedBy=default.target
EOF

    echo -e "${GREEN}✅ Systemd service installed${NC}"
    echo ""
    echo "To enable and start:"
    echo "  systemctl --user daemon-reload"
    echo "  systemctl --user enable qwen-token-refresh"
    echo "  systemctl --user start qwen-token-refresh"
    echo ""
    echo "To check status:"
    echo "  systemctl --user status qwen-token-refresh"
}

# Main
check_dependencies

case "${1:-}" in
    --status|-s)
        show_status
        ;;
    --daemon|-d)
        run_daemon
        ;;
    --install|-i)
        install_systemd_service
        ;;
    --help|-h)
        echo "Usage: $0 [OPTION]"
        echo ""
        echo "Options:"
        echo "  (none)      Check and refresh token if expired"
        echo "  --status    Show token status"
        echo "  --daemon    Run as background refresh daemon"
        echo "  --install   Install as systemd user service"
        echo "  --help      Show this help"
        echo ""
        echo "Examples:"
        echo "  $0                    # Refresh if needed"
        echo "  $0 --status           # Check status"
        echo "  $0 --daemon           # Run background daemon"
        echo "  $0 --install          # Install systemd service"
        ;;
    *)
        # Default: check and refresh if needed
        if [ ! -f "$CREDS_FILE" ]; then
            echo -e "${RED}Error: Credentials not found${NC}"
            echo "Run: qwen --auth-type qwen-oauth -p 'test'"
            exit 1
        fi

        local status=$(get_token_status)

        case "$status" in
            VALID)
                local expiry=$(cat "$CREDS_FILE" | jq -r '.expiry_date // 0')
                echo -e "${GREEN}✅ Token is valid${NC} ($(format_time_remaining $expiry) remaining)"
                echo ""
                echo "To force refresh anyway, delete ~/.qwen/oauth_creds.json and re-auth"
                ;;
            EXPIRED|EXPIRING_SOON)
                echo -e "${YELLOW}Token needs refresh ($status)${NC}"
                refresh_token
                ;;
            *)
                echo -e "${YELLOW}Unknown token status, attempting refresh...${NC}"
                refresh_token
                ;;
        esac
        ;;
esac