<#
.SYNOPSIS
    Flash an ESP-IDF application image to a connected ESP32 board.

.DESCRIPTION
    Auto-detects the serial port, the esptool executable and the flash offset
    (by reading the partition table off the device), then writes the image.

.EXAMPLE
    .\flash.ps1                        # single .bin in this folder, auto everything
    .\flash.ps1 esp_tag.bin -Monitor   # flash then open the serial monitor
    .\flash.ps1 fw.bin 0x20000         # force the flash address (2nd positional arg)
    .\flash.ps1 fw.bin -Port COM12     # force a port
    .\flash.ps1 fw.bin -DryRun         # resolve everything, print the command, write nothing
    .\flash.ps1 -List                  # just list candidate serial ports
#>
[CmdletBinding()]
param(
    [Parameter(Position = 0)]
    [string]$Bin,
    [Parameter(Position = 1)]
    [string]$Offset,
    [string]$Port,
    [int]$Baud = 115200,   # native USB CDC: real speed is USB full-speed; changing baud is flaky on some hosts — keep 115200
    [string]$Esptool,
    [switch]$Monitor,
    [switch]$List,
    [switch]$DryRun,
    [switch]$SkipChipCheck
)

$ErrorActionPreference = 'Stop'
$ScriptDir = Split-Path -Parent $PSCommandPath

$ChipIds = @{
    0  = 'ESP32'
    2  = 'ESP32-S2'
    5  = 'ESP32-C3'
    9  = 'ESP32-S3'
    12 = 'ESP32-C2'
    13 = 'ESP32-C6'
    16 = 'ESP32-H2'
    18 = 'ESP32-P4'
}

$UsbVendors = @{
    '303A' = 'Espressif USB-Serial/JTAG'
    '10C4' = 'Silicon Labs CP210x'
    '1A86' = 'WCH CH340/CH9102'
    '0403' = 'FTDI'
    '067B' = 'Prolific PL2303'
}

function Fail([string]$msg) {
    Write-Host "ERROR: $msg" -ForegroundColor Red
    exit 1
}

function Info([string]$msg) { Write-Host $msg -ForegroundColor Cyan }
function Warn([string]$msg) { Write-Host "WARN: $msg" -ForegroundColor Yellow }

function Read-CString([byte[]]$bytes, [int]$off, [int]$len) {
    $end = $off
    $limit = [Math]::Min($off + $len, $bytes.Length)
    while ($end -lt $limit -and $bytes[$end] -ne 0) { $end++ }
    if ($end -le $off) { return '' }
    return [System.Text.Encoding]::UTF8.GetString($bytes, $off, $end - $off)
}

# ---------------------------------------------------------------- esptool ----
function Find-Esptool {
    if ($Esptool) {
        if (-not (Test-Path $Esptool)) { Fail "esptool not found at: $Esptool" }
        return (Resolve-Path $Esptool).Path
    }
    if ($env:ESPTOOL -and (Test-Path $env:ESPTOOL)) { return $env:ESPTOOL }

    $cmd = Get-Command 'esptool.exe', 'esptool.py', 'esptool' -ErrorAction SilentlyContinue |
           Select-Object -First 1
    if ($cmd) { return $cmd.Source }

    $found = @()
    foreach ($root in @('D:\Espressif', 'C:\Espressif', 'E:\Espressif', "$env:USERPROFILE\.espressif")) {
        if (-not (Test-Path $root)) { continue }
        $found += Get-ChildItem -Path (Join-Path $root 'python_env\*\Scripts\esptool.exe') -ErrorAction SilentlyContinue
        if ($found.Count -eq 0) {
            $found += Get-ChildItem -Path $root -Filter 'esptool.exe' -Recurse -Depth 5 -ErrorAction SilentlyContinue
        }
    }
    if ($found.Count -eq 0) {
        Fail "esptool not found. Install ESP-IDF, or 'pip install esptool', or pass -Esptool <path>."
    }
    # newest-looking environment first (idf5.3... sorts above idf4.4...)
    return ($found | Sort-Object FullName -Descending | Select-Object -First 1).FullName
}

# ------------------------------------------------------------------- port ----
function Get-CandidatePorts {
    $all = @(Get-CimInstance Win32_PnPEntity -ErrorAction SilentlyContinue |
             Where-Object { $_.Name -match '\(COM\d+\)' })
    $out = @()
    foreach ($dev in $all) {
        $com = ''
        if ($dev.Name -match '\((COM\d+)\)') { $com = $Matches[1] }
        $vid = ''
        if ($dev.DeviceID -match 'VID_([0-9A-Fa-f]{4})') { $vid = $Matches[1].ToUpper() }
        $kind = 'other'
        $likely = $false
        if ($UsbVendors.ContainsKey($vid)) { $kind = $UsbVendors[$vid]; $likely = $true }
        # SEGGER J-Link CDC ports are never the ESP target
        if ($vid -eq '1366') { $kind = 'SEGGER J-Link (ignored)'; $likely = $false }
        $out += [pscustomobject]@{
            Port   = $com
            Vid    = $vid
            Kind   = $kind
            Likely = $likely
            Name   = $dev.Name
        }
    }
    return $out | Sort-Object { [int]($_.Port -replace '\D', '') }
}

function Show-Ports($ports) {
    if ($ports.Count -eq 0) { Write-Host '  (no serial ports found)'; return }
    foreach ($p in $ports) {
        $mark = '   '
        if ($p.Likely) { $mark = ' * ' }
        Write-Host ("{0}{1,-6} VID_{2,-5} {3}" -f $mark, $p.Port, $p.Vid, $p.Kind)
    }
}

function Resolve-Port {
    $ports = @(Get-CandidatePorts)
    $likely = @($ports | Where-Object { $_.Likely })

    if ($likely.Count -eq 1) {
        Info ("Port      : {0}  ({1})" -f $likely[0].Port, $likely[0].Kind)
        return $likely[0].Port
    }
    if ($likely.Count -eq 0) {
        Write-Host 'Serial ports seen:'
        Show-Ports $ports
        Fail 'No USB-serial device that looks like an ESP board. Plug it in, or pass -Port COMx.'
    }

    Write-Host 'Multiple candidate ports:'
    Show-Ports $likely
    $ans = Read-Host 'Which port? (e.g. COM8)'
    if (-not $ans) { Fail 'No port chosen.' }
    return $ans.Trim().ToUpper()
}

# ------------------------------------------------------------------ image ----
function Read-ImageInfo([string]$path) {
    $b = [System.IO.File]::ReadAllBytes($path)
    if ($b.Length -lt 0x100) { Fail "$path is too small to be an ESP image." }
    if ($b[0] -ne 0xE9) {
        Fail ("$path does not start with the ESP image magic 0xE9 (found 0x{0:X2})." -f $b[0])
    }
    $chipId = [BitConverter]::ToUInt16($b, 12)
    $chipName = 'unknown'
    if ($ChipIds.ContainsKey([int]$chipId)) { $chipName = $ChipIds[[int]$chipId] }

    $info = [ordered]@{
        Size     = $b.Length
        ChipId   = $chipId
        ChipName = $chipName
        Project  = ''
        Version  = ''
        Built    = ''
        IdfVer   = ''
        IsApp    = $false
    }
    # esp_app_desc_t sits right after the 24-byte header + 8-byte segment header.
    # Magic 0xABCD5432, little endian. Compared byte-wise because PowerShell 5.1
    # parses an 8-digit hex literal as a (negative) Int32.
    if ($b[0x20] -eq 0x32 -and $b[0x21] -eq 0x54 -and $b[0x22] -eq 0xCD -and $b[0x23] -eq 0xAB) {
        $info.IsApp   = $true
        $info.Version = Read-CString $b 0x30 32
        $info.Project = Read-CString $b 0x50 32
        $t = Read-CString $b 0x70 16
        $d = Read-CString $b 0x80 16
        $info.Built   = ("{0} {1}" -f $d, $t).Trim()
        $info.IdfVer  = Read-CString $b 0x90 32
    }
    return [pscustomobject]$info
}

# -------------------------------------------------------- partition table ----
function Get-Partitions([string]$tablePath) {
    $b = [System.IO.File]::ReadAllBytes($tablePath)
    $parts = @()
    for ($i = 0; ($i + 32) -le $b.Length; $i += 32) {
        # 0xAA 0x50 = ESP_PARTITION_MAGIC (0x50AA, little endian)
        if ($b[$i] -ne 0xAA -or $b[$i + 1] -ne 0x50) { break }
        $parts += [pscustomobject]@{
            Type    = $b[$i + 2]
            SubType = $b[$i + 3]
            Offset  = [BitConverter]::ToUInt32($b, $i + 4)
            Size    = [BitConverter]::ToUInt32($b, $i + 8)
            Label   = Read-CString $b ($i + 12) 16
        }
    }
    return $parts
}

function Select-AppPartition($parts) {
    $apps = @($parts | Where-Object { $_.Type -eq 0 })
    if ($apps.Count -eq 0) { return $null }
    $factory = @($apps | Where-Object { $_.SubType -eq 0 })
    if ($factory.Count -eq 1) { return $factory[0] }
    if ($apps.Count -eq 1) { return $apps[0] }
    return $null   # ambiguous (OTA slots): make the user choose with -Offset
}

# ==================================================================== main ====

if ($List) {
    Write-Host 'Serial ports ( * = looks like an ESP board ):'
    Show-Ports (Get-CandidatePorts)
    exit 0
}

# --- resolve the .bin -----------------------------------------------------
if (-not $Bin) {
    $bins = @(Get-ChildItem -Path (Join-Path $ScriptDir '*.bin') -File -ErrorAction SilentlyContinue)
    if ($bins.Count -eq 1) {
        $Bin = $bins[0].FullName
    } elseif ($bins.Count -eq 0) {
        Fail "No .bin given and none found in $ScriptDir. Usage: .\flash.ps1 <file.bin>"
    } else {
        Write-Host 'Several .bin files here - pick one:'
        $bins | ForEach-Object { Write-Host ("  {0}" -f $_.Name) }
        Fail 'Pass the file explicitly, e.g. .\flash.ps1 esp_tag.bin'
    }
}
if (-not (Test-Path $Bin)) {
    $alt = Join-Path $ScriptDir $Bin
    if (Test-Path $alt) { $Bin = $alt } else { Fail "File not found: $Bin" }
}
$Bin = (Resolve-Path $Bin).Path

$img = Read-ImageInfo $Bin
Info ("Image     : {0}" -f (Split-Path -Leaf $Bin))
Write-Host ("            {0} bytes, target {1}" -f $img.Size, $img.ChipName)
if ($img.IsApp) {
    Write-Host ("            project '{0}'  version '{1}'" -f $img.Project, $img.Version)
    Write-Host ("            built {0} with IDF {1}" -f $img.Built, $img.IdfVer)
} else {
    Warn 'No app descriptor found - this may be a bootloader or a merged full-flash image.'
    if (-not $Offset) { Fail 'Not an application image: pass -Offset explicitly (e.g. -Offset 0x0).' }
}

# --- tools and port -------------------------------------------------------
$tool = Find-Esptool
Info ("esptool   : {0}" -f $tool)
if ($Port) { Info ("Port      : {0}  (forced)" -f $Port) } else { $Port = Resolve-Port }

# --- offset: read the device's own partition table ------------------------
if ($Offset) {
    $hex = $Offset -replace '^0[xX]', ''
    try { $addr = [Convert]::ToInt64($hex, 16) }
    catch { Fail "Bad -Offset '$Offset' (expected hex like 0x20000)." }
    Info ("Offset    : 0x{0:X}  (forced)" -f $addr)
} else {
    $tmp = Join-Path $env:TEMP ("ptable_{0}.bin" -f [guid]::NewGuid().ToString('N'))
    Info 'Reading partition table from the device ...'
    $out = & $tool -p $Port read_flash 0x8000 0xC00 $tmp
    if ($LASTEXITCODE -ne 0) { Fail "esptool could not talk to the board on $Port." }

    $joined = ($out | Out-String)
    $detected = ''
    if ($joined -match 'Detecting chip type\.\.\.\s*(\S+)') { $detected = $Matches[1] }
    elseif ($joined -match 'Chip is (\S+)') { $detected = $Matches[1] }
    if ($detected) { Write-Host ("Device    : {0}" -f $detected) }

    if (-not $SkipChipCheck -and $detected -and $img.ChipName -ne 'unknown') {
        if (-not $detected.StartsWith($img.ChipName)) {
            Remove-Item $tmp -ErrorAction SilentlyContinue
            Fail ("Image is built for {0} but the board is {1}. Use -SkipChipCheck to override." -f $img.ChipName, $detected)
        }
    }

    $parts = Get-Partitions $tmp
    Remove-Item $tmp -ErrorAction SilentlyContinue
    if ($parts.Count -eq 0) {
        Fail 'No valid partition table at 0x8000. The board needs a bootloader + partition table first, or pass -Offset.'
    }

    Write-Host 'Partitions:'
    foreach ($p in $parts) {
        Write-Host ("  {0,-16} type {1,-3} sub 0x{2:X2}  0x{3:X6}  {4,8} KB" -f $p.Label, $p.Type, $p.SubType, $p.Offset, [int]($p.Size / 1024))
    }

    $appPart = Select-AppPartition $parts
    if (-not $appPart) {
        Fail 'Could not pick a single app partition (multiple OTA slots?). Pass -Offset explicitly.'
    }
    $addr = $appPart.Offset
    Info ("Offset    : 0x{0:X}  (partition '{1}')" -f $addr, $appPart.Label)

    if ($img.Size -gt $appPart.Size) {
        Fail ("Image is {0} bytes but partition '{1}' is only {2} bytes." -f $img.Size, $appPart.Label, $appPart.Size)
    }
}

# --- flash ----------------------------------------------------------------
$addrArg = ('0x{0:X}' -f $addr)

Write-Host ''
Write-Host 'Command:' -ForegroundColor DarkGray
Write-Host ('  & "{0}" -p {1} -b {2} write_flash {3} "{4}"' -f $tool, $Port, $Baud, $addrArg, $Bin)
Write-Host ''

if ($DryRun) {
    Info 'Dry run - nothing was written.'
    exit 0
}

Info ("Flashing at {0} @ {1} baud ..." -f $addrArg, $Baud)
& $tool -p $Port -b $Baud write_flash $addrArg $Bin
if ($LASTEXITCODE -ne 0) { Fail "esptool write_flash failed (exit $LASTEXITCODE)." }

Write-Host ''
Write-Host 'Flash OK.' -ForegroundColor Green

# --- monitor --------------------------------------------------------------
if ($Monitor) {
    # python.exe lives next to esptool.exe in a venv's Scripts dir; fall back to
    # the env root (some layouts) and finally to whatever is on PATH.
    $toolDir = Split-Path -Parent $tool
    $py = ''
    foreach ($cand in @((Join-Path $toolDir 'python.exe'),
                        (Join-Path (Split-Path -Parent $toolDir) 'python.exe'))) {
        if (Test-Path $cand) { $py = $cand; break }
    }
    if (-not $py) {
        $pyCmd = Get-Command python -ErrorAction SilentlyContinue
        if ($pyCmd) { $py = $pyCmd.Source }
    }
    if (-not $py) {
        Warn 'No python found for the serial monitor; skipping.'
        exit 0
    }
    Write-Host ''
    Info "Monitor on $Port @ 115200 - press Ctrl+] to quit"
    & $py -m serial.tools.miniterm $Port 115200
}