# M5SHARK R112 one-command flasher.
#
#   .\flash-full-fast.ps1                full flash (4 images, first install)
#   .\flash-full-fast.ps1 -Quick         update only (app image, device runs M5SHARK)
#   .\flash-full-fast.ps1 -Port COM7     pick another port
#
# Runs esptool through `python -m` (no PATH problems) and installs it if
# missing. The chip stays in the loader until the run finishes, and a
# read-back verification that is cut off by the USB port AFTER a completed
# write is reported as "verify inconclusive" (press RESET and check the
# screen) instead of a scary failure - the write itself already finished.
param(
    [string]$Port = "COM6",
    [switch]$Quick,
    # Drop to 460800 if the port keeps cutting out mid-write.
    [int]$Baud = 921600
)
$ErrorActionPreference = "Stop"

# Firmware images live next to this script, or one level up in the repo.
$here = Split-Path -Parent $MyInvocation.MyCommand.Path
$binDir = $here
if (-not (Test-Path (Join-Path $binDir "firmware.bin"))) {
    $binDir = Join-Path $here "firmware\USB_FLASH_FILES"
}
foreach ($f in @("bootloader.bin", "partition-table.bin",
                 "ota_data_initial.bin", "firmware.bin")) {
    $p = Join-Path $binDir $f
    if (-not (Test-Path $p)) { throw "missing $f (looked in $binDir)" }
}

# esptool via python -m; install the proven line if absent.
python -m pip show esptool *> $null
if ($LASTEXITCODE -ne 0) {
    Write-Host "installing esptool (once)..."
    python -m pip install esptool
    if ($LASTEXITCODE -ne 0) { throw "pip install failed - is Python installed?" }
}

$common = @("-m", "esptool", "--chip", "esp32c5", "--port", $Port,
            "--baud", "$Baud", "--after", "no_reset")

# Native-command stderr would trip $ErrorActionPreference=Stop while it is
# being captured, so relax it for the esptool run itself.
$prevEap = $ErrorActionPreference
$ErrorActionPreference = "Continue"
if ($Quick) {
    Write-Host "== QUICK UPDATE: firmware.bin @ 0x10000 ==" -ForegroundColor Cyan
    $output = python @common write_flash 0x10000 (Join-Path $binDir "firmware.bin") 2>&1
} else {
    Write-Host "== FULL FLASH: bootloader + partitions + ota + firmware ==" -ForegroundColor Cyan
    $output = python @common write_flash `
        0x2000  (Join-Path $binDir "bootloader.bin") `
        0x8000  (Join-Path $binDir "partition-table.bin") `
        0xE000  (Join-Path $binDir "ota_data_initial.bin") `
        0x10000 (Join-Path $binDir "firmware.bin") 2>&1
}
$ErrorActionPreference = $prevEap
$output | ForEach-Object { Write-Host $_ }
$joined = ($output | Out-String)

if ($LASTEXITCODE -eq 0) {
    Write-Host ""
    Write-Host "Hash of data verified. Press RESET on the device now." -ForegroundColor Green
    Write-Host "Welcome to M5SHARK R112 IQ FAMILY WATCH."
    return
}

# Only a write that FULLY finished (esptool printed its "Wrote ... at
# 0x00010000" line for the firmware and moved on to "Verifying written
# data") may be called complete when the USB drops the read-back. A match
# on any "Wrote N bytes" line is NOT enough - the small images print those
# too, which once misreported a 40%-interrupted write as a success.
$firmwareFullyWritten = $joined -match "(?s)Wrote \d+ bytes[^\r\n]*at 0x00010000[^\r\n]*\r?\n.*Verifying written data"
if ($firmwareFullyWritten -and
    $joined -match "No more data to read|Invalid head of packet") {
    Write-Host ""
    Write-Warning "WRITE COMPLETED - read-back verify was cut off by the USB port."
    Write-Host "Press RESET on the device now and check the boot screen." -ForegroundColor Green
    Write-Host "If it does not come up, hold BOOT, tap RESET, release, and re-run." -ForegroundColor Yellow
    return
}

if ($joined -match "No more data to read|Invalid head of packet") {
    throw "the WRITE was interrupted before the firmware finished - the device does NOT have a complete image. Hold BOOT, tap RESET, release after 2 s and re-run. If it keeps cutting out, re-run with -Baud 460800."
}

throw "flash failed - hold BOOT, tap RESET, release after 2 s and retry"