#!/usr/bin/env bash
set -euo pipefail

############################################################
# WARNING: Do NOT pipe this script to head, tail, less, etc!
# Example: ./flash_cricket_full.sh | head -40   <--- BAD
# This will kill the script mid-flash and leave your device erased!
# Always run this script directly in a terminal.
############################################################

# Full flash Cricket ESP32-P4: erase + bootloader + partitions + app
# Use this for fresh devices or to fix corrupted flash
# Auto-detects directory structure and works with either:
#   1. Full ESP-IDF project directory (uses idf.py flash)
#   2. Binary files only (uses esptool.py directly)
#
# Usage: 
#   flash_cricket_full.sh [port]              # Auto-detect mode
#   flash_cricket_full.sh --build=R006        # Flash from builds/R006/
#   flash_cricket_full.sh --binary            # Force binary mode
#   flash_cricket_full.sh --build=R006 /dev/cu.usbmodem1101
#
# ═══════════════════════════════════════════════════════════════════════════════
# FOR BINARY-ONLY UPGRADES (e.g., R004 → R006):
# ═══════════════════════════════════════════════════════════════════════════════
#
# 1. Download the 3 binaries from builds/R006/:
#      - bootloader.bin
#      - partition-table.bin
#      - cricket.bin
#
# 2. Place them in your cricket directory:
#      uhuman: /Users/uhuman/PycharmProjects/PlatformIO/Projects/M5Tab5-HalUD/platforms/Cricket_ESP32_P4/
#      ming:   /Users/ming/Development/esp32/cricket/
#
#    Files can be in root or in build/ subdirectory:
#      Option A (flat):     cricket.bin, bootloader.bin, partition-table.bin
#      Option B (build/):   build/cricket.bin, build/bootloader.bin, build/partition-table.bin
#
# 3. Run this script:
#      ./flash_cricket_full.sh
#
# 4. After flashing, upload WiFi config:
#      ./flash_wifi_config_merge.sh
#
# ═══════════════════════════════════════════════════════════════════════════════

DEFAULT_PORT="/dev/cu.usbmodem1101"
PORT=""
FORCE_BINARY=0
BUILD_VERSION=""

# Get script directory BEFORE any cd commands
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BUILDS_DIR="$(dirname "$SCRIPT_DIR")/builds"

# Parse arguments
for arg in "$@"; do
  case $arg in
    --binary)
      FORCE_BINARY=1
      ;;
    --build=*)
      BUILD_VERSION="${arg#*=}"
      FORCE_BINARY=1
      ;;
    /dev/*)
      PORT="$arg"
      ;;
    *)
      # Assume it's a port if starts with /
      if [[ "$arg" == /* ]]; then
        PORT="$arg"
      fi
      ;;
  esac
done

PORT="${PORT:-$DEFAULT_PORT}"

echo "=== Cricket ESP32-P4 FULL Flash ==="

# Auto-detect Cricket directory based on user
if [[ "$USER" == "uhuman" ]]; then
  CRICKET_DIR="/Users/uhuman/PycharmProjects/PlatformIO/Projects/M5Tab5-HalUD/platforms/Cricket_ESP32_P4"
elif [[ "$USER" == "ming" ]]; then
  CRICKET_DIR="/Users/ming/Development/esp32/cricket"
else
  echo "ERROR: Unknown user '$USER' - add your path to this script" >&2
  exit 1
fi
echo "Using: $CRICKET_DIR"

# Source ESP-IDF
if [[ -f "$HOME/export-esp.sh" ]]; then
  source "$HOME/export-esp.sh" > /dev/null 2>&1
elif [[ -f "$HOME/esp/esp-idf/export.sh" ]]; then
  source "$HOME/esp/esp-idf/export.sh" > /dev/null 2>&1
else
  echo "ERROR: ESP-IDF not found" >&2
  echo "  Expected: ~/export-esp.sh or ~/esp/esp-idf/export.sh"
  exit 1
fi

cd "$CRICKET_DIR" || { echo "ERROR: Cannot cd to: $CRICKET_DIR" >&2; exit 1; }

# Auto-detect port if default not found
if [[ ! -e "$PORT" ]]; then
  echo "Port $PORT not found, auto-detecting..."
  DETECTED=$(ls /dev/cu.usbmodem* 2>/dev/null | head -1)
  if [[ -n "$DETECTED" ]]; then
    PORT="$DETECTED"
    echo "Using: $PORT"
  else
    echo "ERROR: No USB modem ports found" >&2
    exit 1
  fi
fi

echo ""
echo "Project dir: $CRICKET_DIR"
echo "Port: $PORT"
echo ""

# If --build=R006 specified, use that build folder
if [[ -n "$BUILD_VERSION" ]]; then
  BUILD_DIR="$BUILDS_DIR/$BUILD_VERSION"
  
  if [[ ! -d "$BUILD_DIR" ]]; then
    echo "ERROR: Build version '$BUILD_VERSION' not found in $BUILDS_DIR" >&2
    echo "Available builds:"
    ls -1 "$BUILDS_DIR" 2>/dev/null | grep -E "^R[0-9]+" || echo "  (none)"
    exit 1
  fi
  
  echo "Using build: $BUILD_VERSION (from $BUILD_DIR)"
  cd "$BUILD_DIR"
  FORCE_BINARY=1
fi

# Check if this is a full ESP-IDF project or just binaries
if [[ $FORCE_BINARY -eq 1 ]]; then
  FLASH_MODE="binary"
elif [[ -f "CMakeLists.txt" && -f "sdkconfig" ]]; then
  FLASH_MODE="idf"
  echo "Mode: Full ESP-IDF project (will use idf.py flash)"
  echo ""
else
  FLASH_MODE="binary"
fi

# Binary mode: find and validate files
if [[ "$FLASH_MODE" == "binary" ]]; then
  echo "Mode: Binary files only"
  echo ""
  
  # Find binary files
  if [[ -f "build/cricket.bin" ]]; then
    BIN_DIR="build"
  elif [[ -f "cricket.bin" ]]; then
    BIN_DIR="."
  else
    echo "❌ ERROR: cricket.bin not found!"
    echo ""
    echo "Expected files in: $(pwd)"
    echo "├── bootloader.bin       (or build/bootloader.bin)"
    echo "├── partition-table.bin  (or build/partition-table.bin)"
    echo "└── cricket.bin          (or build/cricket.bin)"
    echo ""
    echo "Download these from builds/R006/ and place them here."
    exit 1
  fi
  
  APP_BIN="$BIN_DIR/cricket.bin"
  BOOTLOADER_BIN="$BIN_DIR/bootloader/bootloader.bin"
  PARTITION_BIN="$BIN_DIR/partition_table/partition-table.bin"
  
  [[ -f "$BOOTLOADER_BIN" ]] || BOOTLOADER_BIN="$BIN_DIR/bootloader.bin"
  [[ -f "$PARTITION_BIN" ]] || PARTITION_BIN="$BIN_DIR/partition-table.bin"
  
  # Show status of each file with date
  echo "Binary files (check dates - all 3 must be updated for upgrade!):"
  echo ""
  for BIN in "$BOOTLOADER_BIN" "$PARTITION_BIN" "$APP_BIN"; do
    if [[ -f "$BIN" ]]; then
      SIZE=$(ls -lh "$BIN" | awk '{print $5}')
      DATE=$(stat -f "%Sm" -t "%Y-%m-%d %H:%M" "$BIN" 2>/dev/null || stat -c "%y" "$BIN" 2>/dev/null | cut -d. -f1)
      echo "  ✅ $(basename "$BIN")"
      echo "     Size: $SIZE  |  Modified: $DATE"
    else
      echo "  ❌ $(basename "$BIN") - NOT FOUND"
    fi
  done
  echo ""
  echo "⚠️  IMPORTANT: All 3 files must be from the SAME build (e.g., R006)!"
  echo "   If dates don't match or are old, download fresh from builds/R006/"
  echo ""
  
  # Check if all required files exist
  MISSING=0
  [[ ! -f "$BOOTLOADER_BIN" ]] && MISSING=1
  [[ ! -f "$PARTITION_BIN" ]] && MISSING=1
  [[ ! -f "$APP_BIN" ]] && MISSING=1
  
  if [[ $MISSING -eq 1 ]]; then
    echo "❌ Some files are missing! Please add them to:"
    echo "   $CRICKET_DIR"
    exit 1
  fi
fi

echo "⚠️  This will ERASE all flash (including WiFi config)"
echo ""
echo "Press ENTER to continue, Ctrl+C to cancel..."
read

echo "Erasing flash..."
esptool.py --chip esp32p4 --port "$PORT" erase_flash

echo ""
echo "Flashing bootloader + partitions + app..."

# Flash based on mode detected earlier
if [[ "$FLASH_MODE" == "idf" ]]; then
  # Full project - use idf.py
  idf.py -p "$PORT" flash
else
  # Binary-only - use esptool.py directly (variables already set above)
  FLASH_CMD="esptool.py --chip esp32p4 --port $PORT --baud 921600"
  FLASH_CMD="$FLASH_CMD --before default_reset --after hard_reset write_flash"
  FLASH_CMD="$FLASH_CMD --flash_mode dio --flash_freq 80m --flash_size 16MB"
  FLASH_CMD="$FLASH_CMD 0x2000 $BOOTLOADER_BIN"
  FLASH_CMD="$FLASH_CMD 0x8000 $PARTITION_BIN"
  FLASH_CMD="$FLASH_CMD 0x10000 $APP_BIN"
  
  eval $FLASH_CMD
fi

echo ""
echo "✅ Full flash complete!"
echo "Note: WiFi config was erased. Run flash_wifi_config_merge.sh to restore."