#!/bin/bash
# ESP32 Mass Deployment Script (Enhanced)
# Features: parallel flashing, auto-retry, verification, progress tracking

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
PROJECT_DIR="${1:-$PWD}"
SSID="${SSID:-ClubHouse2G}"
PASSWORD="${PASSWORD:-}"
START_IP="${START_IP:-192.168.1.51}"
NETMASK="${NETMASK:-255.255.255.0}"
GATEWAY="${GATEWAY:-192.168.1.1}"
BAUD_RATE="${BAUD_RATE:-460800}"
MAX_RETRIES="${MAX_RETRIES:-2}"
VERIFY_PING="${VERIFY_PING:-true}"

# Parse starting IP
IFS='.' read -r -a IP_PARTS <<< "$START_IP"
IP_BASE="${IP_PARTS[0]}.${IP_PARTS[1]}.${IP_PARTS[2]}"
IP_START=${IP_PARTS[3]}

# Check for password
if [ -z "$PASSWORD" ]; then
    echo -e "${RED}ERROR: WiFi password not set!${NC}"
    echo "Usage: PASSWORD='your_wifi_pass' $0 [project_dir]"
    echo ""
    echo "Or set via environment:"
    echo "  export PASSWORD='your_wifi_pass'"
    echo "  export SSID='YourSSID'          # Default: ClubHouse2G"
    echo "  export START_IP='192.168.1.51'  # Default: 192.168.1.51"
    echo "  $0"
    exit 1
fi

print_banner() {
    echo ""
    echo -e "${BLUE}=========================================="
    echo "ESP32 Mass Deployment Tool"
    echo -e "==========================================${NC}"
    echo "Project:  $PROJECT_DIR"
    echo "SSID:     $SSID"
    echo "IP Range: ${IP_BASE}.${IP_START}+"
    echo "Gateway:  $GATEWAY"
    echo "Netmask:  $NETMASK"
    echo "Verify:   $VERIFY_PING"
    echo -e "${BLUE}==========================================${NC}"
}

print_banner

# Step 1: Build firmware
echo ""
echo -e "${YELLOW}[1/4] Building firmware...${NC}"
cd "$PROJECT_DIR"
if ! idf.py build; then
    echo -e "${RED}✗ Build failed!${NC}"
    exit 1
fi
echo -e "${GREEN}✓ Build complete${NC}"

# Step 2: Detect devices
echo ""
echo -e "${YELLOW}[2/4] Detecting ESP32 devices...${NC}"
DEVICES=($(ls /dev/ttyUSB* /dev/ttyACM* 2>/dev/null || true))

if [ ${#DEVICES[@]} -eq 0 ]; then
    echo -e "${RED}ERROR: No devices found!${NC}"
    echo "Connect ESP32 devices via USB and try again."
    exit 1
fi

echo -e "${GREEN}Found ${#DEVICES[@]} device(s):${NC}"
for i in "${!DEVICES[@]}"; do
    IP_ADDR="${IP_BASE}.$((IP_START + i))"
    echo "  [$i] ${DEVICES[$i]} → $IP_ADDR"
done

# Step 3: Flash and configure
echo ""
echo -e "${YELLOW}[3/4] Flashing and configuring...${NC}"
echo ""

flash_and_configure() {
    local INDEX=$1
    local DEVICE=$2
    local IP_ADDR="${IP_BASE}.$((IP_START + INDEX))"
    local LOG_FILE="/tmp/esp32_deploy_${INDEX}.log"
    local STATUS_FILE="/tmp/esp32_status_${INDEX}"
    
    {
        for ATTEMPT in $(seq 1 $MAX_RETRIES); do
            echo "=== Device $INDEX: $DEVICE (Attempt $ATTEMPT/$MAX_RETRIES) ==="
            echo "Target IP: $IP_ADDR"
            
            # Flash
            echo "Flashing..."
            if idf.py -p "$DEVICE" -b "$BAUD_RATE" flash 2>&1; then
                echo "✓ Flash successful"
            else
                echo "✗ Flash failed on attempt $ATTEMPT"
                if [ $ATTEMPT -eq $MAX_RETRIES ]; then
                    echo "FAILED" > "$STATUS_FILE"
                    exit 1
                fi
                sleep 2
                continue
            fi
            
            # Wait for boot
            sleep 3
            
            # Configure WiFi
            echo "Configuring WiFi..."
            {
                echo "CFG"
                echo "SSID:$SSID"
                echo "PASS:$PASSWORD"
                echo "IP:$IP_ADDR"
                echo "MASK:$NETMASK"
                echo "GW:$GATEWAY"
                echo "DHCP:0"
                echo "END"
            } > "$DEVICE" 2>/dev/null || true
            
            # Wait for network
            sleep 5
            
            # Verify if requested
            if [ "$VERIFY_PING" = "true" ]; then
                echo "Verifying connectivity..."
                if ping -c 2 -W 3 "$IP_ADDR" > /dev/null 2>&1; then
                    echo "✓ Ping successful"
                    echo "SUCCESS" > "$STATUS_FILE"
                    break
                else
                    echo "✗ Ping failed on attempt $ATTEMPT"
                    if [ $ATTEMPT -eq $MAX_RETRIES ]; then
                        echo "PING_FAIL" > "$STATUS_FILE"
                    fi
                fi
            else
                echo "SUCCESS" > "$STATUS_FILE"
                break
            fi
            
            sleep 2
        done
        
    } > "$LOG_FILE" 2>&1
    
    # Show result
    if [ -f "$STATUS_FILE" ]; then
        STATUS=$(cat "$STATUS_FILE")
        if [ "$STATUS" = "SUCCESS" ]; then
            echo -e "${GREEN}[Device $INDEX] ${DEVICES[$INDEX]} → $IP_ADDR [OK]${NC}"
        elif [ "$STATUS" = "PING_FAIL" ]; then
            echo -e "${YELLOW}[Device $INDEX] ${DEVICES[$INDEX]} → $IP_ADDR [FLASHED, NO PING]${NC}"
        else
            echo -e "${RED}[Device $INDEX] ${DEVICES[$INDEX]} → $IP_ADDR [FAILED]${NC}"
        fi
    fi
}

export -f flash_and_configure
export PROJECT_DIR SSID PASSWORD IP_BASE IP_START NETMASK GATEWAY BAUD_RATE MAX_RETRIES VERIFY_PING
export RED GREEN YELLOW BLUE NC

# Clean old status files
rm -f /tmp/esp32_status_* /tmp/esp32_deploy_*.log

# Launch parallel jobs
for i in "${!DEVICES[@]}"; do
    flash_and_configure "$i" "${DEVICES[$i]}" &
done

# Wait for completion
wait

# Step 4: Summary
echo ""
echo -e "${YELLOW}[4/4] Deployment Summary${NC}"
echo -e "${BLUE}==========================================${NC}"

SUCCESS_COUNT=0
FAILED_COUNT=0
PING_FAIL_COUNT=0

for i in "${!DEVICES[@]}"; do
    IP_ADDR="${IP_BASE}.$((IP_START + i))"
    STATUS_FILE="/tmp/esp32_status_${i}"
    
    if [ -f "$STATUS_FILE" ]; then
        STATUS=$(cat "$STATUS_FILE")
        case "$STATUS" in
            SUCCESS)
                echo -e "${GREEN}✓${NC} ${DEVICES[$i]} → $IP_ADDR"
                ((SUCCESS_COUNT++))
                ;;
            PING_FAIL)
                echo -e "${YELLOW}⚠${NC} ${DEVICES[$i]} → $IP_ADDR (no ping response)"
                ((PING_FAIL_COUNT++))
                ;;
            *)
                echo -e "${RED}✗${NC} ${DEVICES[$i]} → $IP_ADDR (failed)"
                ((FAILED_COUNT++))
                ;;
        esac
    else
        echo -e "${RED}✗${NC} ${DEVICES[$i]} → $IP_ADDR (no status)"
        ((FAILED_COUNT++))
    fi
done

echo -e "${BLUE}==========================================${NC}"
echo -e "Total:   ${#DEVICES[@]} devices"
echo -e "${GREEN}Success: $SUCCESS_COUNT${NC}"
if [ $PING_FAIL_COUNT -gt 0 ]; then
    echo -e "${YELLOW}Warning: $PING_FAIL_COUNT (flashed but no ping)${NC}"
fi
if [ $FAILED_COUNT -gt 0 ]; then
    echo -e "${RED}Failed:  $FAILED_COUNT${NC}"
fi
echo -e "${BLUE}==========================================${NC}"

echo ""
echo "Logs: /tmp/esp32_deploy_*.log"
echo ""
echo "Test commands:"
echo "  iperf -c ${IP_BASE}.${IP_START}"
echo "  for i in {${IP_START}..$((IP_START + ${#DEVICES[@]} - 1))}; do ping -c 1 ${IP_BASE}.\$i & done; wait"
echo ""

# Cleanup status files
rm -f /tmp/esp32_status_*

exit $FAILED_COUNT