#!/bin/bash
# Flash wifi_config.txt to Cricket's LittleFS partition via USB-C
# MERGE version - reads existing LittleFS, merges WiFi config, preserves all files
# This is the SAFE version that won't delete your recordings!

set -e  # Exit on error

# Detect platform-specific config file locations
CONFIG_FOUND=""
if [ -f "/Users/uhuman/PycharmProjects/wifi_config.txt" ]; then
    DEFAULT_CONFIG="/Users/uhuman/PycharmProjects/wifi_config.txt"
    CONFIG_FOUND="uhuman PycharmProjects"
elif [ -f "/Users/ming/Development/PlatformIO/Projects/M5Tab5-HalUD/platforms/Cricket_ESP32_P4/wifi_config.txt" ]; then
    DEFAULT_CONFIG="/Users/ming/Development/PlatformIO/Projects/M5Tab5-HalUD/platforms/Cricket_ESP32_P4/wifi_config.txt"
    CONFIG_FOUND="ming PlatformIO"
elif [ -f "./wifi_config.txt" ]; then
    DEFAULT_CONFIG="./wifi_config.txt"
    CONFIG_FOUND="current directory"
else
    DEFAULT_CONFIG=""
fi

CONFIG_FILE="${1:-$DEFAULT_CONFIG}"
DEFAULT_PORT="/dev/cu.usbmodem1101"
PORT="${2:-$DEFAULT_PORT}"

echo "========================================"
echo "Flash WiFi Config to Cricket (MERGE)"
echo "========================================"
echo ""

# Check if config file exists
if [ -z "$CONFIG_FILE" ] || [ ! -f "$CONFIG_FILE" ]; then
    echo "❌ Error: wifi_config.txt not found!"
    echo ""
    echo "The script looks for wifi_config.txt in these locations (in order):"
    echo "  1. /Users/uhuman/PycharmProjects/wifi_config.txt"
    echo "  2. /Users/ming/Development/esp32/cricket/wifi_config.txt"
    echo "  3. ./wifi_config.txt (current directory)"
    echo ""
    echo "Please create wifi_config.txt in one of these locations with format:"
    echo "  SSID=your_network_name"
    echo "  PASSWORD=your_password"
    echo ""
    echo "Or specify the path manually:"
    echo "  $0 /path/to/wifi_config.txt"
    exit 1
fi

echo "✓ Config file found in: $CONFIG_FOUND"
echo "📝 Config file: $CONFIG_FILE"
echo "🔌 USB port: $PORT"
echo ""

# Show credentials
echo "WiFi credentials to flash:"
cat "$CONFIG_FILE"
echo ""

# Validate format
if ! grep -q "^SSID=" "$CONFIG_FILE" || ! grep -q "^PASSWORD=" "$CONFIG_FILE"; then
    echo "❌ Error: Invalid wifi_config.txt format!"
    echo ""
    echo "Required format:"
    echo "  SSID=your_network_name"
    echo "  PASSWORD=your_password"
    echo ""
    echo "Each entry must be on its own line with no extra spaces around the = sign."
    exit 1
fi

# Auto-detect ESP-IDF
if [ -f "$HOME/export-esp.sh" ]; then
    echo "Loading ESP-IDF from $HOME/export-esp.sh..."
    source "$HOME/export-esp.sh"
elif [ -f "$HOME/esp/esp-idf/export.sh" ]; then
    echo "Loading ESP-IDF from $HOME/esp/esp-idf/export.sh..."
    source "$HOME/esp/esp-idf/export.sh"
else
    echo "❌ Error: ESP-IDF not found."
    echo "   Install ESP-IDF first: ./install_esp_idf.sh"
    exit 1
fi

# Check if port exists, auto-detect if not
if [ ! -e "$PORT" ]; then
    echo "⚠️  Port $PORT not found, auto-detecting..."
    DETECTED_PORT=$(ls /dev/cu.usbmodem* 2>/dev/null | head -n 1)
    if [ -z "$DETECTED_PORT" ]; then
        echo "❌ No USB device found. Check:"
        echo "   - Cricket is connected via USB-C"
        echo "   - USB cable supports data (not charge-only)"
        exit 1
    fi
    PORT="$DETECTED_PORT"
    echo "✓ Auto-detected port: $PORT"
fi

# Check if mklittlefs is available
if ! command -v mklittlefs &> /dev/null; then
    echo "❌ Error: mklittlefs not found."
    echo ""
    echo "Install mklittlefs (one-time setup):"
    echo "  brew install mklittlefs"
    echo ""
    echo "Or download from:"
    echo "  https://github.com/earlephilhower/mklittlefs/releases"
    echo ""
    echo "For macOS ARM64, download 'arm64-apple-darwin-mklittlefs'"
    echo "Then: chmod +x mklittlefs && sudo mv mklittlefs /usr/local/bin/"
    exit 1
fi

echo ""
echo "🔄 MERGE MODE: Existing files will be preserved!"
echo ""
echo "This script will:"
echo "  1. Read existing LittleFS partition from Cricket"
echo "  2. Extract all files (recordings, sniffer audio, etc.)"
echo "  3. Add/update wifi_config.txt with new credentials"
echo "  4. Write merged filesystem back to Cricket"
echo ""
echo "Press ENTER to continue, or Ctrl+C to cancel..."
read

# Create temporary directory
TEMP_DIR=$(mktemp -d)
trap "rm -rf $TEMP_DIR" EXIT

# LittleFS partition offset from partitions.csv:
# factory(10MB) @ 0x10000 + coredump(64K) @ 0xA10000 + storage(1MB) @ 0xA20000 = littlefs @ 0xB20000
LITTLEFS_OFFSET="0xB20000"
LITTLEFS_SIZE="4194304"  # 4MB
LITTLEFS_BACKUP="$TEMP_DIR/littlefs_backup.bin"
LITTLEFS_EXTRACT="$TEMP_DIR/fs_extract"
LITTLEFS_NEW="$TEMP_DIR/littlefs_new.bin"

echo ""
echo "📥 Step 1/4: Reading existing LittleFS from Cricket..."
esptool.py --chip esp32p4 \
    --port "$PORT" \
    --baud 921600 \
    read_flash \
    "$LITTLEFS_OFFSET" "$LITTLEFS_SIZE" \
    "$LITTLEFS_BACKUP"

if [ $? -ne 0 ]; then
    echo "❌ Failed to read LittleFS partition. Check USB connection."
    exit 1
fi

echo ""
echo "📂 Step 2/4: Extracting existing files..."
mkdir -p "$LITTLEFS_EXTRACT"

# Try to unpack existing filesystem (may fail if empty/corrupt - that's OK)
mklittlefs -u "$LITTLEFS_EXTRACT" -b 4096 -p 256 -s "$LITTLEFS_SIZE" "$LITTLEFS_BACKUP" 2>/dev/null || {
    echo "⚠️  No existing filesystem found (fresh device) - will create new one"
}

# List existing files
if [ "$(ls -A $LITTLEFS_EXTRACT)" ]; then
    echo "✓ Found existing files:"
    ls -lh "$LITTLEFS_EXTRACT"
else
    echo "✓ No existing files (fresh device)"
fi

echo ""
echo "📝 Step 3/4: Merging WiFi config..."
cp "$CONFIG_FILE" "$LITTLEFS_EXTRACT/wifi_config.txt"
echo "✓ WiFi config added/updated in filesystem"

echo ""
echo "📦 Step 4/4: Creating merged LittleFS image..."
mklittlefs -c "$LITTLEFS_EXTRACT" -b 4096 -p 256 -s "$LITTLEFS_SIZE" "$LITTLEFS_NEW"

echo ""
echo "⚡ Flashing merged LittleFS partition..."
esptool.py --chip esp32p4 \
    --port "$PORT" \
    --baud 921600 \
    write_flash \
    "$LITTLEFS_OFFSET" "$LITTLEFS_NEW"

if [ $? -eq 0 ]; then
    echo ""
    echo "✅ Success! WiFi config merged into Cricket's LittleFS."
    echo ""
    
    # Verify by reading back and checking
    echo "🔍 Verifying flash..."
    VERIFY_IMG="$TEMP_DIR/verify.bin"
    esptool.py --chip esp32p4 \
        --port "$PORT" \
        --baud 921600 \
        read_flash \
        "$LITTLEFS_OFFSET" "$LITTLEFS_SIZE" \
        "$VERIFY_IMG" > /dev/null 2>&1
    
    if [ $? -eq 0 ]; then
        # Try to extract and verify wifi_config.txt exists
        VERIFY_DIR="$TEMP_DIR/verify_extract"
        mkdir -p "$VERIFY_DIR"
        mklittlefs -u "$VERIFY_DIR" -b 4096 -p 256 -s "$LITTLEFS_SIZE" "$VERIFY_IMG" 2>/dev/null
        
        if [ -f "$VERIFY_DIR/wifi_config.txt" ]; then
            echo "✓ Verified: wifi_config.txt successfully written to Cricket!"
            echo "✓ Contents:"
            cat "$VERIFY_DIR/wifi_config.txt"
        else
            echo "⚠️  Warning: Could not verify wifi_config.txt (but flash reported success)"
        fi
    fi
    
    echo ""
    echo "All existing files preserved:"
    ls -lh "$LITTLEFS_EXTRACT" | grep -v "^total" | grep -v "^d"
    echo ""
    echo "Next steps:"
    echo "  1. Power cycle Cricket (unplug and replug USB)"
    echo "  2. Check serial console - should see: '✓ WiFi config loaded from /littlefs/wifi_config.txt'"
    echo "  3. Cricket will connect to your WiFi automatically"
    echo ""
    echo "💡 Your recordings and sniffer audio are safe!"
else
    echo ""
    echo "❌ Flash failed. Your original data is still intact on the device."
    echo "   Check USB connection and try again."
    exit 1
fi