#!/usr/bin/env bash
# ============================================================
# HF Account Management Tool
#
# Support multi-user authorization switching
#
# Usage:
#   hf-account.sh add <username> [token]     Add/update account
#   hf-account.sh list                       List all accounts
#   hf-account.sh use <username>             Switch default account
#   hf-account.sh current                   Show current account
#   hf-account.sh remove <username>         Delete account
#   hf-account.sh get-token [username]       Get account token
#
# Example:
#   hf-account.sh add user1 hf_xxxxx        Add account
#   hf-account.sh add user2                 Interactive add
#   hf-account.sh list                      List all accounts
#   hf-account.sh use user1                 Switch to user1
#   hf-account.sh remove user2              Delete account
#   hf-account.sh current                   Show current account
#   hf-account.sh get-token user1            Get user1's token
#
# Config file:
#   ~/.config/openclaw/hf-accounts.json
#
# ============================================================

set -euo pipefail

# ---- Fix terminal for interactive input ----
# Handle backspace and special characters properly
if [[ -t 0 ]]; then
    stty sane 2>/dev/null || true
fi

# ---- Find Python with huggingface_hub ----
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
VENV_DIR="$SCRIPT_DIR/../.venv"

# Use venv python directly if it exists (cross-platform: Windows uses Scripts/, POSIX uses bin/)
VENV_BIN_DIR=""
if [[ -f "$VENV_DIR/Scripts/python.exe" ]]; then
    PYTHON_CMD="$VENV_DIR/Scripts/python.exe"
    VENV_BIN_DIR="$VENV_DIR/Scripts"
elif [[ -f "$VENV_DIR/bin/python" ]]; then
    PYTHON_CMD="$VENV_DIR/bin/python"
    VENV_BIN_DIR="$VENV_DIR/bin"
elif command -v python3 >/dev/null 2>&1; then
    PYTHON_CMD="python3"
else
    PYTHON_CMD="python"
fi

if [[ -n "$VENV_BIN_DIR" ]]; then
    export PATH="$VENV_BIN_DIR:$PATH"
fi

# ---- Config ----
HF_ACCOUNTS_FILE="${HF_ACCOUNTS_FILE:-$HOME/.config/openclaw/hf-accounts.json}"
HF_TOKEN_FILE="${HF_TOKEN_FILE:-$HOME/.cache/huggingface/token}"

# ---- Colors ----
# 使用 $'...' 让 bash 解释 \033 为 ESC 字节，这样在 cat/printf/echo 中都能正确着色
RED=$'\033[0;31m'; GREEN=$'\033[0;32m'; YELLOW=$'\033[1;33m'
CYAN=$'\033[0;36m'; BOLD=$'\033[1m'; NC=$'\033[0m'

info()    { printf "${GREEN}[INFO]${NC} %s\n" "$*"; }
warn()    { printf "${YELLOW}[WARN]${NC} %s\n" "$*" >&2; }
error()   { printf "${RED}[ERROR]${NC} %s\n" "$*" >&2; }
success() { printf "${CYAN}[OK]${NC} %s\n" "$*"; }

# ---- Ensure config dir exists ----
ensure_config_dir() {
    local config_dir
    config_dir="$(dirname "$HF_ACCOUNTS_FILE")"
    if [[ ! -d "$config_dir" ]]; then
        mkdir -p "$config_dir"
    fi
}

# ---- Init config file ----
init_accounts_file() {
    ensure_config_dir
    if [[ ! -f "$HF_ACCOUNTS_FILE" ]]; then
        cat > "$HF_ACCOUNTS_FILE" <<'EOF'
{
  "accounts": {},
  "current": null
}
EOF
    fi
}

# ---- Read accounts ----
read_accounts() {
    init_accounts_file
    cat "$HF_ACCOUNTS_FILE"
}

# ---- Save accounts ----
save_accounts() {
    local data="$1"
    ensure_config_dir
    echo "$data" | $PYTHON_CMD -c "
import json, sys
data = json.load(sys.stdin)
with open('$HF_ACCOUNTS_FILE', 'w') as f:
    json.dump(data, f, indent=2)
"
}

# ---- Get current account name ----
get_current_account() {
    local data
    data="$(read_accounts 2>/dev/null || echo '{}')"
    echo "$data" | $PYTHON_CMD -c "
import json, sys
try:
    data = json.loads(sys.stdin.read() or '{}')
except:
    data = {}
current = data.get('current', '')
print(current if current else '')
"
}

# ---- Get account token ----
get_account_token() {
    local account="$1"
    local data
    data="$(read_accounts 2>/dev/null || echo '{}')"
    echo "$data" | $PYTHON_CMD -c "
import json, sys
try:
    data = json.loads(sys.stdin.read() or '{}')
except:
    data = {}
accounts = data.get('accounts', {})
token = accounts.get('$account', '')
print(token if token else '')
"
}

# ---- Check if account exists ----
account_exists() {
    local account="$1"
    local data
    data="$(read_accounts 2>/dev/null || echo '{}')"
    local exists
    exists=$(echo "$data" | $PYTHON_CMD -c "
import json, sys
try:
    data = json.loads(sys.stdin.read() or '{}')
except:
    data = {}
accounts = data.get('accounts', {})
print('yes' if '$account' in accounts else 'no')
")
    [[ "$exists" == "yes" ]]
}

# ---- Interactive account picker ----
# Lists accounts (current first, then alphabetical) and prompts user
# to pick a number. Echoes the chosen username on stdout.
# Returns: 0 + name on stdout = success; 1 = cancel / no accounts; 2 = non-TTY.
pick_account() {
    local prompt="${1:-选择账号}"

    if [[ ! -t 0 ]]; then
        return 2
    fi

    local data
    data="$(read_accounts 2>/dev/null || echo '{}')"

    # Build numbered list: current first, then alphabetical.
    # Output lines: "<index>\t<name>\t<marker>" (marker may be empty).
    local entries
    entries=$(echo "$data" | $PYTHON_CMD -c "
import json, sys
try:
    data = json.loads(sys.stdin.read() or '{}')
except:
    data = {}
names = sorted((data.get('accounts') or {}).keys())
current = data.get('current', '') or ''
ordered = []
if current in names:
    ordered.append(current)
    names.remove(current)
ordered.extend(names)
for i, name in enumerate(ordered, 1):
    marker = ' [current]' if name == current else ''
    print(f'{i}\t{name}\t{marker}')
")

    if [[ -z "$entries" ]]; then
        warn "没有可用账号，请先使用 '$0 add <user> <token>' 添加"
        return 1
    fi

    # All UI goes to stderr so stdout only carries the chosen name
    # (callers do `var=$(pick_account ...)`).
    echo "$prompt:" >&2
    local max_idx=0
    while IFS=$'\t' read -r idx name marker; do
        [[ -z "$idx" ]] && continue
        printf "  %d) %s%s\n" "$idx" "$name" "$marker" >&2
        (( idx > max_idx )) && max_idx=$idx
    done <<< "$entries"
    printf "  9) 取消\n" >&2

    while true; do
        read -r -p "请选择 [1-$max_idx, 9]: " choice >&2
        case "$choice" in
            9|"") return 1 ;;
            *)
                if [[ "$choice" =~ ^[0-9]+$ ]] && (( choice >= 1 && choice <= max_idx )); then
                    echo "$entries" | awk -F'\t' -v c="$choice" '$1 == c { print $2; exit }'
                    return 0
                fi
                warn "无效选择"
                ;;
        esac
    done
}

# ---- Add account ----
cmd_add() {
    local account="${1:-}"
    local token="${2:-}"

    if [[ -z "$account" ]]; then
        error "Username cannot be empty"
        echo "Usage: hf-account.sh add <username> [token]"
        exit 1
    fi

    # If token not provided, read interactively
    if [[ -z "$token" ]]; then
        if [[ -t 0 ]]; then
            # TTY: silent input (no echo to terminal)
            read -rs -p "Enter HF Token (input hidden, press Enter): " token
            echo
        else
            # Non-TTY: read with echo (e.g. piped input or CI); warn the user
            warn "Reading HF_TOKEN from stdin (will be visible in process list)"
            read -r token
        fi
        if [[ -z "$token" ]]; then
            error "Token cannot be empty"
            exit 1
        fi
    fi

    # Verify token
    info "Verifying Token..."
    local valid_username
    local verify_result
    verify_result=$(HF_TOKEN="$token" $PYTHON_CMD -c "
import os
import sys
try:
    from huggingface_hub import HfApi
except ImportError as e:
    print(f'ERROR:huggingface_hub_not_installed:{e}')
    sys.exit(1)
token = os.environ.get('HF_TOKEN', '').strip()
if not token:
    print('ERROR:empty_token')
    sys.exit(1)
api = HfApi(token=token)
try:
    data = api.whoami()
    name = data.get('name', '')
    if name:
        print(name)
    else:
        print('ERROR:no_name_returned')
except Exception as e:
    print(f'ERROR:{type(e).__name__}:{e}')
" 2>&1) || true

    # Check if result starts with ERROR:
    if [[ "$verify_result" == ERROR:* ]]; then
        error "Token verification failed: ${verify_result#ERROR:}"
        exit 1
    fi

    valid_username="$verify_result"

    if [[ -z "$valid_username" ]]; then
        error "Token verification failed: empty response"
        exit 1
    fi

    # Check if token username matches
    if [[ "$valid_username" != "$account" ]]; then
        warn "Token corresponds to user '$valid_username', not '$account'"
        if [[ -t 0 ]]; then
            read -r -p "Continue using '$account' as account name? (y/n): " confirm
            [[ "$confirm" != "y" ]] && info "Cancelled" && exit 0
        else
            error "Cannot confirm account name mismatch in non-interactive mode."
            exit 1
        fi
    fi

    # Read existing data
    local data
    data="$(read_accounts)"

    # Add/update account
    data=$(echo "$data" | HF_TOKEN="$token" $PYTHON_CMD -c "
import os, json, sys
data = json.load(sys.stdin)
token = os.environ.get('HF_TOKEN', '')
data['accounts']['$account'] = token
if data.get('current') is None:
    data['current'] = '$account'
print(json.dumps(data, indent=2))
")

    save_accounts "$data"
    success "Account '$account' added/updated"

    # Sync to default token file (current account)
    local current
    current=$(get_current_account)
    if [[ "$current" == "$account" ]]; then
        sync_to_default_token
    fi
}

# ---- List all accounts ----
cmd_list() {
    local data
    data="$(read_accounts 2>/dev/null || echo '{"accounts": {}, "current": null}')"

    local current
    current=$(get_current_account 2>/dev/null || echo "")

    echo ""
    echo "HF Account List:"
    echo ""

    echo "$data" | $PYTHON_CMD -c "
import json, sys

try:
    raw = sys.stdin.read()
    if not raw or not raw.strip():
        data = {'accounts': {}, 'current': ''}
    else:
        data = json.loads(raw)
except (json.JSONDecodeError, ValueError):
    data = {'accounts': {}, 'current': ''}

accounts = data.get('accounts', {})
current = data.get('current', '') or ''

if not accounts:
    print('  (none)')
else:
    for name, token in sorted(accounts.items()):
        marker = ''
        if name == current:
            marker = ' [current]'
        print(f'  {name}{marker}')
        print(f'    Token: {token}')
        print()

print(f'Current: {current}')
"

    echo ""
    echo "Current account: ${current:-none}"
    echo ""
}

# ---- Switch account ----
cmd_use() {
    local account="${1:-}"

    if [[ -z "$account" ]] && [[ -t 0 ]]; then
        # Interactive: show picker. Any picker failure (cancel / no accounts)
        # exits 0 — the user just backed out. The [[ -t 0 ]] guard means
        # pick_account's "non-TTY" return (2) cannot fire here; in non-TTY
        # mode we fall through to the "Username cannot be empty" error below.
        account=$(pick_account "选择要切换的账号") || exit 0
    fi

    if [[ -z "$account" ]]; then
        error "Username cannot be empty"
        echo "Usage: hf-account.sh use <username>"
        exit 1
    fi

    if ! account_exists "$account"; then
        error "Account '$account' does not exist"
        echo "Use 'hf-account.sh list' to view all accounts"
        exit 1
    fi

    # Update current account
    local data
    data="$(read_accounts)"
    data=$(echo "$data" | $PYTHON_CMD -c "
import json, sys
data = json.load(sys.stdin)
data['current'] = '$account'
print(json.dumps(data, indent=2))
")
    save_accounts "$data"

    # Sync token to default file
    sync_to_default_token

    success "Switched to account: $account"
}

# ---- Show current account ----
cmd_current() {
    local current
    current=$(get_current_account)

    if [[ -z "$current" ]]; then
        echo "No current account"
        exit 0
    fi

    echo "$current"
}

# ---- Delete account ----
cmd_remove() {
    local account="${1:-}"

    if [[ -z "$account" ]]; then
        error "Username cannot be empty"
        echo "Usage: hf-account.sh remove <username>"
        exit 1
    fi

    if ! account_exists "$account"; then
        error "Account '$account' does not exist"
        exit 1
    fi

    # Read data
    local data
    data="$(read_accounts)"
    local current
    current=$(get_current_account)

    # If deleting current account
    if [[ "$current" == "$account" ]]; then
        warn "About to delete current account"
        # Find another account as new current
        local new_current
        new_current=$(echo "$data" | $PYTHON_CMD -c "
import json, sys
data = json.load(sys.stdin)
accounts = data.get('accounts', {})
for name in accounts:
    if name != '$account':
        print(name)
        break
")
        data=$(echo "$data" | $PYTHON_CMD -c "
import json, sys
data = json.load(sys.stdin)
data['current'] = '${new_current:-null}'
print(json.dumps(data, indent=2))
")
    fi

    # Delete account
    data=$(echo "$data" | $PYTHON_CMD -c "
import json, sys
data = json.load(sys.stdin)
if '$account' in data.get('accounts', {}):
    del data['accounts']['$account']
print(json.dumps(data, indent=2))
")
    save_accounts "$data"

    # If deleted current account, sync new token
    if [[ "$current" == "$account" && -n "$new_current" ]]; then
        sync_to_default_token
    fi

    success "Account '$account' deleted"
}

# ---- Sync from default token file ----
cmd_sync_from_default() {
    local default_token_file="${HF_TOKEN_FILE:-$HOME/.cache/huggingface/token}"

    if [[ ! -f "$default_token_file" ]]; then
        error "Default token file not found: $default_token_file"
        echo "Please ensure you have logged in via 'huggingface-cli login' or add an account first"
        exit 1
    fi

    info "Reading token from $default_token_file..."
    local token
    token="$(cat "$default_token_file" 2>/dev/null || echo "")"

    if [[ -z "$token" ]]; then
        error "Token file is empty"
        exit 1
    fi

    # Verify token and get username
    info "Verifying token..."
    local valid_username
    local verify_result
    verify_result=$(HF_TOKEN="$token" $PYTHON_CMD -c "
import os
import sys
try:
    from huggingface_hub import HfApi
except ImportError as e:
    print(f'ERROR:huggingface_hub_not_installed:{e}')
    sys.exit(1)
token = os.environ.get('HF_TOKEN', '').strip()
api = HfApi(token=token)
try:
    data = api.whoami()
    name = data.get('name', '')
    if name:
        print(name)
    else:
        print('ERROR:no_name_returned')
except Exception as e:
    print(f'ERROR:{type(e).__name__}:{e}')
" 2>&1) || true

    if [[ "$verify_result" == ERROR:* ]]; then
        error "Token verification failed: ${verify_result#ERROR:}"
        exit 1
    fi

    valid_username="$verify_result"

    if [[ -z "$valid_username" ]]; then
        error "Token verification failed: empty response"
        exit 1
    fi

    info "Token belongs to user: $valid_username"

    # Check if account already exists
    if account_exists "$valid_username"; then
        local current
        current=$(get_current_account)
        if [[ "$current" == "$valid_username" ]]; then
            info "Account '$valid_username' already exists and is current"
        else
            warn "Account '$valid_username' already exists"
            if [[ -t 0 ]]; then
                read -r -p "Switch to this account? (y/n): " confirm
                if [[ "$confirm" == "y" ]]; then
                    cmd_use "$valid_username"
                fi
            else
                error "Cannot confirm switch in non-interactive mode."
                exit 1
            fi
        fi
    else
        # Auto-create account
        if [[ -t 0 ]]; then
            read -r -p "Create new account '$valid_username' and set as current? (y/n): " confirm
            if [[ "$confirm" != "y" ]]; then
                info "Cancelled"
                exit 0
            fi
        else
            error "Cannot confirm account creation in non-interactive mode."
            exit 1
        fi

        # Add account with the token
        local data
        data="$(read_accounts)"
        data=$(echo "$data" | HF_TOKEN="$token" $PYTHON_CMD -c "
import os, json, sys
data = json.load(sys.stdin)
token = os.environ.get('HF_TOKEN', '')
data['accounts']['$valid_username'] = token
data['current'] = '$valid_username'
print(json.dumps(data, indent=2))
")
        save_accounts "$data"
        success "Account '$valid_username' created and set as current"
    fi
}

# ---- Get token ----
cmd_get_token() {
    local account="${1:-}"

    # If not specified, use current account
    if [[ -z "$account" ]]; then
        account=$(get_current_account)
    fi

    if [[ -z "$account" ]]; then
        error "No current account, add account first or specify account name"
        exit 1
    fi

    local token
    token=$(get_account_token "$account")

    if [[ -z "$token" ]]; then
        error "Account '$account' does not exist or has no token"
        exit 1
    fi

    printf '%s' "$token"
}

# ---- Sync current account to default token file ----
sync_to_default_token() {
    local current
    current=$(get_current_account)

    if [[ -z "$current" ]]; then
        return 1
    fi

    local token
    token=$(get_account_token "$current")

    if [[ -n "$token" ]]; then
        ensure_config_dir
        echo "$token" > "$HF_TOKEN_FILE"
        chmod 600 "$HF_TOKEN_FILE"
    fi
}

# ---- Sync to environment variable ----
export_current_token() {
    local current
    current=$(get_current_account)

    if [[ -z "$current" ]]; then
        return 1
    fi

    local token
    token=$(get_account_token "$current")

    if [[ -n "$token" ]]; then
        export HF_TOKEN="$token"
    fi
}

# ---- Help ----
usage() {
    cat <<EOF
${BOLD}HF Account Management Tool${NC}

${BOLD}Usage:${NC}
    $0 <command> [args...]

${BOLD}Commands:${NC}
    add <username> [token]     Add/update account (token optional, interactive input)
    list                       List all accounts (full tokens shown for copy-paste)
    use [username]             Switch default account (no arg: interactive picker)
    current                    Show current account
    remove <username>          Delete account
    get-token [username]       Get account token (default current account)
    sync-from-default          Sync from default token file (~/.cache/huggingface/token)

${BOLD}Examples:${NC}
    $0 add user1 hf_xxxxx         # Add account
    $0 add user2                  # Interactive add
    $0 list                       # List all accounts
    $0 use                        # Switch via interactive picker
    $0 use user1                  # Switch to user1 (script-friendly)
    $0 remove user2               # Delete account
    $0 current                    # Show current account
    $0 get-token user1            # Get user1's token
    $0 sync-from-default          # Import from ~/.cache/huggingface/token

${BOLD}Config file:${NC}
    $HF_ACCOUNTS_FILE

${BOLD}Features:${NC}
    - Multi-account: Add multiple HF accounts and switch between them
    - Interactive picker: 'use' (no arg) shows numbered account list
    - Current account: All HF tools use current account's token by default
    - Auto-sync: Switching accounts automatically updates ~/.cache/huggingface/token
    - Token security: Tokens are stored in config file (list prints full token for copy-paste)

EOF
}

# ============================================================
# Main entry
# ============================================================

main() {
    if [[ $# -eq 0 ]]; then
        usage
        exit 0
    fi

    local cmd="$1"
    shift

    case "$cmd" in
        add)
            cmd_add "$@"
            ;;
        list)
            cmd_list
            ;;
        use)
            cmd_use "$@"
            ;;
        current)
            cmd_current
            ;;
        remove|rm)
            cmd_remove "$@"
            ;;
        get-token)
            cmd_get_token "$@"
            ;;
        sync-from-default|sync)
            cmd_sync_from_default
            ;;
        -h|--help|help)
            usage
            ;;
        *)
            error "Unknown command: $cmd"
            usage
            exit 1
            ;;
    esac
}

main "$@"