#!/usr/bin/env bash

PYTHON=${PYTHON:-$(which python3 python | head -n 1)}
BPS_RESET=false
MCU=""

# Constants
RESET_BAUD=1200
FIRMWARE_OFFSET=0x00
# Fallback offsets, used only when firmware metadata carries no partition table.
# These track partition-table.csv; boards with their own table override them below.
SPIFFS_OFFSET=0x300000
OTA_OFFSET=0x260000

# Determine the correct esptool command to use
if "$PYTHON" -m esptool version >/dev/null 2>&1; then
    ESPTOOL_CMD="$PYTHON -m esptool"
elif command -v esptool >/dev/null 2>&1; then
    ESPTOOL_CMD="esptool"
elif command -v esptool.py >/dev/null 2>&1; then
    ESPTOOL_CMD="esptool.py"
else
    echo "Error: esptool not found"
    exit 1
fi

# Check for jq
if ! command -v jq >/dev/null 2>&1; then
    echo "Error: jq not found" >&2
    echo "Install jq with your package manager." >&2
    echo "e.g. 'apt install jq', 'dnf install jq', 'brew install jq', etc." >&2
    exit 1
fi

# esptool v5 supports commands with dashes and deprecates commands with
# underscores. Prior versions only support commands with underscores
if ${ESPTOOL_CMD} | grep --quiet write-flash
then
    ESPTOOL_WRITE_FLASH=write-flash
    ESPTOOL_ERASE_FLASH=erase-flash
    ESPTOOL_READ_FLASH_STATUS=read-flash-status
else
    ESPTOOL_WRITE_FLASH=write_flash
    ESPTOOL_ERASE_FLASH=erase_flash
    ESPTOOL_READ_FLASH_STATUS=read_flash_status
fi

set -e

# Usage info
show_help() {
    cat <<EOF
Usage: $(basename "$0") [-h] [-p ESPTOOL_PORT] [-P PYTHON] [-f FILENAME] [--1200bps-reset]
Flash image file to device, but first erasing and writing system information.

    -h               Display this help and exit.
    -p ESPTOOL_PORT  Set the environment variable for ESPTOOL_PORT.  If not set, ESPTOOL iterates all ports (Dangerous).
    -P PYTHON        Specify alternate python interpreter to use to invoke esptool. (Default: "$PYTHON")
    -f FILENAME      The firmware *.factory.bin file to flash.  Custom to your device type and region.
    --1200bps-reset  Attempt to place the device in correct mode. Some hardware requires this twice. (1200bps Reset)

EOF
}
# Parse arguments using a single while loop
while [ $# -gt 0 ]; do
    case "$1" in
    -h | --help)
        show_help
        exit 0
        ;;
    -p)
        ESPTOOL_CMD="$ESPTOOL_CMD --port $2"
        shift
        ;;
    -P)
        PYTHON="$2"
        shift
        ;;
    -f)
        FILENAME="$2"
        shift
        ;;
    --1200bps-reset)
        BPS_RESET=true
        ;;
    --) # Stop parsing options
        shift
        break
        ;;
    *)
        echo "Unknown argument: $1" >&2
        exit 1
        ;;
    esac
    shift # Move to the next argument
done

if [[ $BPS_RESET == true ]]; then
    $ESPTOOL_CMD --baud $RESET_BAUD --after no_reset ${ESPTOOL_READ_FLASH_STATUS}
    exit 0
fi

[ -z "$FILENAME" ] && [ -n "$1" ] && {
    FILENAME="$1"
    shift
}

if [[ $(basename "$FILENAME") != firmware-*.factory.bin ]]; then
  echo "Filename must be a firmware-*.factory.bin file."
  exit 1
fi

# Extract PROGNAME from %FILENAME% for later use.
PROGNAME="${FILENAME/.factory.bin/}"
# Derive metadata filename from %PROGNAME%.
METAFILE="${PROGNAME}.mt.json"

if [[ -f "$FILENAME" && "$FILENAME" == *.factory.bin ]]; then
    # Display metadata if it exists
    if [[ -f "$METAFILE" ]]; then
        echo "Firmware metadata: ${METAFILE}"
        jq . "$METAFILE"
        # Extract relevant fields from metadata
        if [[ $(jq -r '.part' "$METAFILE") != "null" ]]; then
            META_OTA_OFFSET=$(jq -r '.part[] | select(.subtype == "ota_1") | .offset' "$METAFILE")
            META_SPIFFS_OFFSET=$(jq -r '.part[] | select(.subtype == "spiffs") | .offset' "$METAFILE")
            # A partition table may omit either entry; keep the built-in default then.
            if [[ -n "$META_OTA_OFFSET" && "$META_OTA_OFFSET" != "null" ]]; then
                OTA_OFFSET="$META_OTA_OFFSET"
            fi
            if [[ -n "$META_SPIFFS_OFFSET" && "$META_SPIFFS_OFFSET" != "null" ]]; then
                SPIFFS_OFFSET="$META_SPIFFS_OFFSET"
            fi
        fi
        MCU=$(jq -r '.mcu' "$METAFILE")
    else
        echo "ERROR: No metadata file found at ${METAFILE}"
        exit 1
    fi

    # Determine OTA filename based on MCU type (unified OTA format)
    OTAFILE="mt-${MCU}-ota.bin"

    # Set SPIFFS filename with "littlefs-" prefix.
    SPIFFSFILE="littlefs-${PROGNAME/firmware-/}.bin"

    if [[ ! -f "$FILENAME" ]]; then
        echo "Error: file ${FILENAME} wasn't found. Terminating."
        exit 1
    fi
    if [[ ! -f "$OTAFILE" ]]; then
        echo "Error: file ${OTAFILE} wasn't found. Terminating."
        exit 1
    fi
    if [[ ! -f "$SPIFFSFILE" ]]; then
        echo "Error: file ${SPIFFSFILE} wasn't found. Terminating."
        exit 1
    fi

    echo "Trying to flash ${FILENAME}, but first erasing and writing system information"
    $ESPTOOL_CMD ${ESPTOOL_ERASE_FLASH}
    $ESPTOOL_CMD ${ESPTOOL_WRITE_FLASH} $FIRMWARE_OFFSET "${FILENAME}"
    echo "Trying to flash ${OTAFILE} at offset ${OTA_OFFSET}"
    $ESPTOOL_CMD ${ESPTOOL_WRITE_FLASH} "$OTA_OFFSET" "${OTAFILE}"
    echo "Trying to flash ${SPIFFSFILE}, at offset ${SPIFFS_OFFSET}"
    $ESPTOOL_CMD ${ESPTOOL_WRITE_FLASH} "$SPIFFS_OFFSET" "${SPIFFSFILE}"

else
    show_help
    echo "Invalid file: ${FILENAME}"
fi

exit 0