#!/usr/bin/env sh
set -eu

# Environment variables:
#   PORT   - serial port (default: /dev/ttyUSB0)
#   BAUD   - baud rate (default: 460800)
#   FIRMWARE - local path to firmware .bin
#   FW_URL  - URL to download firmware if FIRMWARE not provided or missing

PORT=${PORT:-}
BAUD=${BAUD:-460800}
CHIP=${CHIP:-esp32} # esp32 | esp32c3

# Map chip-specific settings
ESPT_CHIP="$CHIP"
FLASH_OFFSET=0x1000
DEFAULT_DOWNLOAD_PAGE="https://micropython.org/download/ESP32/"
BOARD_ID="ESP32_GENERIC"
BOARD_PAGE="https://micropython.org/download/${BOARD_ID}/"
case "$CHIP" in
  esp32c3)
    ESPT_CHIP="esp32c3"
    FLASH_OFFSET=0x0
    DEFAULT_DOWNLOAD_PAGE="https://micropython.org/download/ESP32C3/"
    BOARD_ID="ESP32_GENERIC_C3"
    BOARD_PAGE="https://micropython.org/download/${BOARD_ID}/"
    ;;
  esp32)
    ESPT_CHIP="esp32"
    FLASH_OFFSET=0x1000
    DEFAULT_DOWNLOAD_PAGE="https://micropython.org/download/ESP32/"
    BOARD_ID="ESP32_GENERIC"
    BOARD_PAGE="https://micropython.org/download/${BOARD_ID}/"
    ;;
  *)
    echo "Unsupported CHIP: $CHIP (supported: esp32, esp32c3)" >&2
    exit 1
    ;;
esac

# Download-only mode: fetch the appropriate firmware and exit
if [ -n "${DOWNLOAD_ONLY:-}" ]; then
  # Prefer resolving latest if nothing provided
  if [ -z "${FIRMWARE:-}" ] && [ -z "${FW_URL:-}" ]; then
    LATEST=1
  fi
  if ! resolve_firmware; then
    echo "Failed to resolve firmware for CHIP=$CHIP" >&2
    exit 1
  fi
  echo "$FIRMWARE"
  exit 0
fi

# Helper: resolve the latest firmware URL for a given board pattern with multiple fallbacks
resolve_latest_url() {
  board_pattern="$1" # e.g., ESP32_GENERIC_C3-.*\.bin
  # Candidate pages to try in order
  pages="${BOARD_PAGE} ${DOWNLOAD_PAGE:-$DEFAULT_DOWNLOAD_PAGE} https://micropython.org/download/ https://micropython.org/resources/firmware/"
  for page in $pages; do
    echo "Trying to resolve latest from $page" >&2
    html=$(curl -fsSL -H 'User-Agent: Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36' -e 'https://micropython.org/download/' "$page" || true)
    [ -z "$html" ] && continue
    # Prefer matching the board pattern
    url=$(printf "%s" "$html" \
      | sed -n 's/.*href=\"\([^\"]*\.bin\)\".*/\1/p' \
      | grep -E "$board_pattern" \
      | head -n1)
    if [ -n "$url" ]; then
      case "$url" in
        http*) echo "$url"; return 0 ;;
        /*) echo "https://micropython.org$url"; return 0 ;;
        *) echo "$page$url"; return 0 ;;
      esac
    fi
  done
  return 1
}

# If LATEST is set and neither FIRMWARE nor FW_URL are provided, auto-detect latest URL
if [ -n "${LATEST:-}" ] && [ -z "${FIRMWARE:-}" ] && [ -z "${FW_URL:-}" ]; then
  # Default board identifiers for each chip
  case "$CHIP" in
    esp32c3) BOARD_ID="ESP32_GENERIC_C3" ;;
    esp32) BOARD_ID="ESP32_GENERIC" ;;
    *) BOARD_ID="ESP32_GENERIC" ;;
  esac
  pattern="${BOARD_ID}-.*\\.bin"
  echo "Resolving latest firmware for $BOARD_ID"
  if FW_URL=$(resolve_latest_url "$pattern"); then
    export FW_URL
    echo "Latest firmware resolved to: $FW_URL"
  else
    echo "Failed to resolve latest firmware for pattern $pattern" >&2
    exit 1
  fi
fi

# Resolve firmware path, downloading if needed
resolve_firmware() {
  if [ -z "${FIRMWARE:-}" ]; then
    if [ -n "${FW_URL:-}" ] || [ -n "${LATEST:-}" ]; then
      # If FW_URL still unset, resolve latest using board-specific pattern
      if [ -z "${FW_URL:-}" ]; then
        case "$CHIP" in
          esp32c3) BOARD_ID="ESP32_GENERIC_C3" ;;
          esp32) BOARD_ID="ESP32_GENERIC" ;;
          *) BOARD_ID="ESP32_GENERIC" ;;
        esac
        pattern="${BOARD_ID}-.*\\.bin"
        echo "Resolving latest firmware for $BOARD_ID"
        if ! FW_URL=$(resolve_latest_url "$pattern"); then
          echo "Failed to resolve latest firmware for pattern $pattern" >&2
          exit 1
        fi
      fi
      mkdir -p .cache
      FIRMWARE=".cache/$(basename "$FW_URL")"
      if [ ! -f "$FIRMWARE" ]; then
        echo "Downloading firmware from $FW_URL to $FIRMWARE"
        curl -L --fail -o "$FIRMWARE" "$FW_URL"
      else
        echo "Firmware already downloaded at $FIRMWARE"
      fi
    else
      # Default fallback: fetch latest using board-specific pattern
      case "$CHIP" in
        esp32c3) BOARD_ID="ESP32_GENERIC_C3" ;;
        esp32) BOARD_ID="ESP32_GENERIC" ;;
        *) BOARD_ID="ESP32_GENERIC" ;;
      esac
      pattern="${BOARD_ID}-.*\\.bin"
      echo "No FIRMWARE or FW_URL specified. Auto-fetching latest for $BOARD_ID"
      if ! FW_URL=$(resolve_latest_url "$pattern"); then
        echo "Failed to resolve latest firmware for pattern $pattern" >&2
        exit 1
      fi
      mkdir -p .cache
      FIRMWARE=".cache/$(basename "$FW_URL")"
      if [ ! -f "$FIRMWARE" ]; then
        echo "Downloading firmware from $FW_URL to $FIRMWARE"
        curl -L --fail -o "$FIRMWARE" "$FW_URL"
      else
        echo "Firmware already downloaded at $FIRMWARE"
      fi
    fi
  else
    if [ ! -f "$FIRMWARE" ]; then
      if [ -n "${FW_URL:-}" ]; then
        mkdir -p "$(dirname "$FIRMWARE")"
        echo "Firmware not found at $FIRMWARE. Downloading from $FW_URL"
        curl -L --fail -o "$FIRMWARE" "$FW_URL"
      else
        echo "Firmware file not found: $FIRMWARE. Provide FW_URL to download automatically." >&2
        exit 1
      fi
    fi
  fi
}

# Auto-detect PORT if not specified
if [ -z "$PORT" ]; then
  candidates="$(ls /dev/tty/ACM* /dev/tty/USB* 2>/dev/null || true)"
  # Some systems expose without /dev/tty/ prefix patterns; try common Linux paths
  [ -z "$candidates" ] && candidates="$(ls /dev/ttyACM* /dev/ttyUSB* 2>/dev/null || true)"
  # Prefer ACM (often for C3) then USB
  PORT=$(printf "%s\n" $candidates | grep -E "/dev/ttyACM[0-9]+" | head -n1 || true)
  [ -z "$PORT" ] && PORT=$(printf "%s\n" $candidates | grep -E "/dev/ttyUSB[0-9]+" | head -n1 || true)
  if [ -z "$PORT" ]; then
    echo "No serial port detected. Connect the board and set PORT=/dev/ttyACM0 (or /dev/ttyUSB0)." >&2
    exit 1
  fi
  echo "Auto-detected PORT=$PORT"
fi

# Preflight: ensure port exists
if [ ! -e "$PORT" ]; then
  echo "Port $PORT does not exist. Detected candidates:" >&2
  ls /dev/ttyACM* /dev/ttyUSB* 2>/dev/null || true
  exit 1
fi

ESPL="python -m esptool"

detect_chip() {
  # Try to detect actual connected chip using esptool and override if needed
  out=$($ESPL --port "$PORT" --baud "$BAUD" chip_id 2>&1 || true)
  case "$out" in
    *"ESP32-C3"*) DETECTED_CHIP=esp32c3 ;;
    *"ESP32"*) DETECTED_CHIP=esp32 ;;
    *) DETECTED_CHIP="" ;;
  esac
  if [ -n "$DETECTED_CHIP" ] && [ "$DETECTED_CHIP" != "$ESPT_CHIP" ]; then
    echo "Detected chip $DETECTED_CHIP differs from requested $ESPT_CHIP. Using detected chip."
    ESPT_CHIP="$DETECTED_CHIP"
    case "$ESPT_CHIP" in
      esp32c3) FLASH_OFFSET=0x0 ;;
      esp32) FLASH_OFFSET=0x1000 ;;
    esac
  fi
}

detect_chip

# Now that we know the actual chip, resolve the correct firmware for it
resolve_firmware

# Validate firmware matches detected chip; if not, auto-correct by fetching the right image
EXPECTED_BOARD_ID="ESP32_GENERIC"
case "$ESPT_CHIP" in
  esp32c3) EXPECTED_BOARD_ID="ESP32_GENERIC_C3" ;;
  esp32) EXPECTED_BOARD_ID="ESP32_GENERIC" ;;

esac

FW_BASENAME="$(basename "$FIRMWARE")"
case "$FW_BASENAME" in
  ${EXPECTED_BOARD_ID}-*.bin) : ;; # ok
  *)
    echo "Firmware $FW_BASENAME does not match detected chip ($ESPT_CHIP). Fetching correct image for $EXPECTED_BOARD_ID..."
    pattern="${EXPECTED_BOARD_ID}-.*\\.bin"
    if ! FW_URL=$(resolve_latest_url "$pattern"); then
      echo "Failed to resolve a firmware matching $EXPECTED_BOARD_ID" >&2
      exit 1
    fi
    mkdir -p .cache
    FIRMWARE=".cache/$(basename "$FW_URL")"
    if [ ! -f "$FIRMWARE" ]; then
      echo "Downloading firmware from $FW_URL to $FIRMWARE"
      curl -L --fail -o "$FIRMWARE" "$FW_URL"
    else
      echo "Firmware already downloaded at $FIRMWARE"
    fi
    ;;
esac

$ESPL --chip "$ESPT_CHIP" --port "$PORT" --baud "$BAUD" erase_flash

echo "Writing firmware $FIRMWARE to $FLASH_OFFSET..."
$ESPL --chip "$ESPT_CHIP" --port "$PORT" --baud "$BAUD" write_flash -z "$FLASH_OFFSET" "$FIRMWARE"

echo "Done."