#!/bin/bash

if ! command -v bluetoothctl &> /dev/null; then
    echo "ERROR: bluetoothctl not found. Please install bluez."
    exit 1
fi

DEVICE_CACHE="/tmp/.bt_dev_cache_$(whoami)"

show_help() {
    cat << 'EOF'
Bluetooth Control Script
Usage: $0 <command> [options]
Commands:
  on         Enable Bluetooth
  off        Disable Bluetooth
  status     Show controller status
  list       Scan 8s and list discovered devices (paired first)
  connect    Connect to device by number or MAC
  disconnect Disconnect device by number or MAC
  pair       Pair device by number or MAC
  remove     Remove unpair device by number or MAC
  help       Show this help
Example:
  $0 list       # Scan and list devices
  $0 connect 1  # Connect to device #1
  $0 pair 1     # Pair and connect to device #1
EOF
}

DEVICES=()
NAMES=()
PAIRED_DEVICES=()
PAIRED_NAMES=()
UNPAIRED_DEVICES=()
UNPAIRED_NAMES=()

build_device_list() {
    PAIRED_DEVICES=()
    PAIRED_NAMES=()
    UNPAIRED_DEVICES=()
    UNPAIRED_NAMES=()
    while read -r line; do
        [[ -z "$line" ]] && continue
        mac=$(echo "$line" | awk '{print $2}')
        name=$(echo "$line" | cut -d' ' -f3-)
        local skip=0
        if [[ "$name" == "$mac" ]]; then skip=1
        elif [[ "$name" =~ ^[0-9A-Fa-f][0-9A-Fa-f](-[0-9A-Fa-f][0-9A-Fa-f]){5}$ ]]; then skip=1; fi
        if [[ $skip -eq 0 ]]; then
            PAIRED_DEVICES+=("$mac")
            PAIRED_NAMES+=("$name")
        fi
    done < <(printf "devices Paired\nquit\n" | bluetoothctl 2>/dev/null | grep '^Device')
    while read -r line; do
        [[ -z "$line" ]] && continue
        mac=$(echo "$line" | awk '{print $2}')
        name=$(echo "$line" | cut -d' ' -f3-)
        local is_paired=0
        for paired_mac in "${PAIRED_DEVICES[@]}"; do
            if [[ "$paired_mac" == "$mac" ]]; then is_paired=1; break; fi
        done
        [[ $is_paired -eq 1 ]] && continue
        local skip=0
        if [[ "$name" == "$mac" ]]; then skip=1
        elif [[ "$name" =~ ^[0-9A-Fa-f][0-9A-Fa-f](-[0-9A-Fa-f][0-9A-Fa-f]){5}$ ]]; then skip=1; fi
        if [[ $skip -eq 0 ]]; then
            UNPAIRED_DEVICES+=("$mac")
            UNPAIRED_NAMES+=("$name")
        fi
    done < <(printf "devices\nquit\n" | bluetoothctl 2>/dev/null | grep '^Device')
    DEVICES=("${PAIRED_DEVICES[@]}" "${UNPAIRED_DEVICES[@]}")
    NAMES=("${PAIRED_NAMES[@]}" "${UNPAIRED_NAMES[@]}")
}

save_device_cache() {
    : > "$DEVICE_CACHE"
    local i=0
    while [[ $i -lt ${#DEVICES[@]} ]]; do
        echo "${DEVICES[$i]}|${NAMES[$i]}" >> "$DEVICE_CACHE"
        ((i++))
    done
}

load_device_cache() {
    DEVICES=()
    NAMES=()
    if [[ ! -f "$DEVICE_CACHE" ]]; then return 1; fi
    while IFS='|' read -r mac name; do
        [[ -z "$mac" ]] && continue
        DEVICES+=("$mac")
        NAMES+=("$name")
    done < "$DEVICE_CACHE"
    [[ ${#DEVICES[@]} -gt 0 ]]
}

is_paired() {
    printf "info %s\nquit\n" "$1" | bluetoothctl 2>/dev/null | grep -qi "Paired: yes"
}

is_connected() {
    printf "info %s\nquit\n" "$1" | bluetoothctl 2>/dev/null | grep -qi "Connected: yes"
}

cmd_on() {
    bluetoothctl power on > /dev/null 2>&1
    bluetoothctl discoverable on > /dev/null 2>&1
    bluetoothctl pairable on > /dev/null 2>&1
    echo "Bluetooth enabled"
}

cmd_off() {
    bluetoothctl power off > /dev/null 2>&1
    echo "Bluetooth disabled"
}

cmd_status() {
    bluetoothctl show
}

cmd_list() {
    bluetoothctl power on > /dev/null 2>&1
    { printf "scan on\n"; sleep 8; printf "scan off\nquit\n"; } | bluetoothctl > /dev/null 2>&1
    echo "Discovered devices (paired first):"
    echo "-----------------------------------"
    build_device_list
    save_device_cache
    
    local paired_count=${#PAIRED_DEVICES[@]}
    local i=1
    while [[ $i -le ${#DEVICES[@]} ]]; do
        local name="${NAMES[$((i-1))]}"
        if [[ $i -le $paired_count ]]; then
            echo "$i. [Paired] $name"
        else
            echo "$i. [New]    $name"
        fi
        ((i++))
    done
    
    local total_unpaired=${#UNPAIRED_DEVICES[@]}
    echo "-----------------------------------"
    echo "Total: ${#DEVICES[@]} devices ($paired_count paired, $total_unpaired new)"
}

get_mac() {
    local input=$1
    if [[ "$input" =~ ^[0-9]+$ ]]; then
        echo "${DEVICES[$((input - 1))]}"
    else
        echo "$input"
    fi
}

cmd_connect() {
    local input=$1
    if [[ -z "$input" ]]; then echo "ERROR: Please provide device number or MAC"; exit 1; fi
    load_device_cache || build_device_list
    local mac=$(get_mac "$input")
    if [[ -z "$mac" || ${#DEVICES[@]} -eq 0 ]]; then echo "ERROR: Invalid device number"; exit 1; fi

    if is_connected "$mac"; then
        echo "✓ Already connected."
        return 0
    fi

    echo "Connecting to $mac..."
    { echo "connect $mac"; sleep 10; } | bluetoothctl > /dev/null 2>&1

    if is_connected "$mac"; then
        echo "✓ Connected successfully!"
    else
        echo "✗ Connection failed. Try again."
    fi
}

cmd_disconnect() {
    local input=$1
    if [[ -z "$input" ]]; then echo "ERROR: Please provide device number or MAC"; exit 1; fi
    load_device_cache || build_device_list
    local mac=$(get_mac "$input")
    if [[ -z "$mac" || ${#DEVICES[@]} -eq 0 ]]; then echo "ERROR: Invalid device number"; exit 1; fi
    echo "Disconnecting from $mac..."
    printf "disconnect %s\nquit\n" "$mac" | bluetoothctl
}

cmd_pair() {
    local input=$1
    if [[ -z "$input" ]]; then echo "ERROR: Please provide device number or MAC"; exit 1; fi

    bluetoothctl power on > /dev/null 2>&1

    # 1. Load MAC from cache (编号不会变)
    if [[ "$input" =~ ^[0-9]+$ ]]; then
        if load_device_cache; then
            echo "Using cached device list."
        else
            echo "No cached list. Please run '$0 list' first."
            exit 1
        fi
    else
        # MAC 输入模式也要扫描
        { printf "scan on\n"; sleep 10; printf "scan off\nquit\n"; } | bluetoothctl > /dev/null 2>&1
        build_device_list
        save_device_cache
    fi

    local mac=$(get_mac "$input")
    local name=""
    if [[ "$input" =~ ^[0-9]+$ ]]; then
        name="${NAMES[$((input - 1))]}"
    fi
    if [[ -z "$mac" || ${#DEVICES[@]} -eq 0 ]]; then
        echo "ERROR: Invalid device number. Please run '$0 list' first."
        exit 1
    fi

    echo "Target: ${name:-$mac} ($mac)"

    # --------------------------------------------------------
    # STEP 1: Check if already paired
    # --------------------------------------------------------
    if is_paired "$mac"; then
        echo "✓ Already paired."
    else
        echo "Starting pairing..."
        echo "Scanning (5s) to make device visible..."

        # 重新扫描让设备回到 bluetoothd 活动列表
        # 注意：不重建设备列表，MAC 仍从缓存取，编号不会变
        { printf "scan on\n"; sleep 5; } | bluetoothctl > /dev/null 2>&1

        # pair + yes，丢弃输出（太嘈杂无法可靠解析）
        {
            echo "pair $mac"
            sleep 5
            echo "yes"
            sleep 3
        } | bluetoothctl > /dev/null 2>&1

        # 停止扫描
        { printf "scan off\n"; sleep 1; } | bluetoothctl > /dev/null 2>&1

        # 等待 bluetoothd 写入配对密钥
        sleep 2

        # 用 info 查询真实状态（不依赖输出解析）
        if is_paired "$mac"; then
            echo "✓ Pairing successful!"
        else
            echo "✗ Pairing failed."
            echo "  Possible reasons:"
            echo "  - Device is not in pairing mode"
            echo "  - Device is out of range"
            echo "  - Try: put device in pairing mode, then run '$0 pair $input'"
            exit 1
        fi
    fi

    # --------------------------------------------------------
    # STEP 2: Wait for device to initialize
    # --------------------------------------------------------
    echo "Waiting for device to initialize (5s)..."
    sleep 5

    # --------------------------------------------------------
    # STEP 3: Check auto-connect (TWS earbuds often do this)
    # --------------------------------------------------------
    if is_connected "$mac"; then
        echo "✓ Connected successfully! (Auto-connected by device)"
        return 0
    fi

    # --------------------------------------------------------
    # STEP 4: Connect with retry
    # --------------------------------------------------------
    echo "Establishing connection..."
    local retries=3
    local attempt=1

    while [[ $attempt -le $retries ]]; do
        echo "  Attempt $attempt/$retries..."

        # 扫描让设备可见
        { printf "scan on\n"; sleep 3; } | bluetoothctl > /dev/null 2>&1

        # 尝试连接
        { echo "connect $mac"; sleep 8; } | bluetoothctl > /dev/null 2>&1

        # 停止扫描
        { printf "scan off\n"; sleep 1; } | bluetoothctl > /dev/null 2>&1

        # 查询真实状态
        if is_connected "$mac"; then
            echo "✓ Connected successfully!"
            return 0
        fi

        if [[ $attempt -lt $retries ]]; then
            echo "  Not connected yet, waiting 4s before retry..."
            sleep 4
        fi
        ((attempt++))
    done

    echo "⚠ Connection failed after $retries attempts."
    echo "  Try: $0 connect $input"
}

cmd_remove() {
    local input=$1
    if [[ -z "$input" ]]; then echo "ERROR: Please provide device number or MAC"; exit 1; fi
    load_device_cache || build_device_list
    local mac=$(get_mac "$input")
    if [[ -z "$mac" || ${#DEVICES[@]} -eq 0 ]]; then echo "ERROR: Invalid device number"; exit 1; fi
    echo "Removing device $mac..."
    printf "remove %s\nquit\n" "$mac" | bluetoothctl
}

case "$1" in
    on)         cmd_on ;;
    off)        cmd_off ;;
    status)     cmd_status ;;
    list)       cmd_list ;;
    connect)    cmd_connect "$2" ;;
    disconnect) cmd_disconnect "$2" ;;
    pair)       cmd_pair "$2" ;;
    remove)     cmd_remove "$2" ;;
    help)       show_help ;;
    "")         show_help ;;
    *)          echo "Unknown command: $1"; show_help; exit 1 ;;
esac

exit 0

