#!/usr/bin/env python3
"""Прошивка RT-4D из WSL через powershell.exe (WSL2 не отдаёт COM напрямую).

    python3 tools/flash_ps.py stock-fw/rt4d_ru_batt.bin [--port COM4]

Рация должна быть в режиме прошивки: выключить -> включить, удерживая тангенту PTT.

Протокол бутлоадера (тот же, что в flash_rt4d.py):
    рукопожатие : слать 0xFF, пока не придёт 0xFF
    стирание    : [39 33 05 10]+CK, затем [39 33 05 55]+CK  -> ACK 0x06
    запись      : [57 hi lo] + 1024 Б + CK                  -> ACK 0x06
    CK = (0x48 + сумма) & 0xFF
"""
import argparse, os, shutil, subprocess, sys

WIN_TMP_WSL = "/mnt/c/Users/vikto/AppData/Local/Temp"
WIN_TMP_WIN = r"C:\Users\vikto\AppData\Local\Temp"

PS = r'''
$ErrorActionPreference = 'Stop'
[Console]::OutputEncoding = [System.Text.Encoding]::UTF8
$fw = [System.IO.File]::ReadAllBytes("{fwpath}")
$p = New-Object System.IO.Ports.SerialPort "{port}",115200,"None",8,"one"
$p.ReadTimeout = 250
$p.WriteTimeout = 5000
try {{ $p.Open() }} catch {{ Write-Output "ERR: порт занят или недоступен: $($_.Exception.Message)"; exit 1 }}

Write-Output "INFO: прошивка $($fw.Length) байт, порт {port}"

# --- рукопожатие ---
$p.DiscardInBuffer()
$ready = $false
$sw = [Diagnostics.Stopwatch]::StartNew()
while ($sw.ElapsedMilliseconds -lt 8000) {{
  try {{ if ($p.ReadByte() -eq 0xFF) {{ $ready = $true; break }} }}
  catch {{ $p.Write([byte[]]@(0xFF), 0, 1) }}
}}
if (-not $ready) {{
  Write-Output "ERR: бутлоадер не отвечает. Выключите рацию и включите с зажатой тангентой PTT."
  $p.Close(); exit 1
}}
Write-Output "INFO: бутлоадер найден"

# --- стирание ---
foreach ($trig in 0x10, 0x55) {{
  $pl = [byte[]]@(0x39, 0x33, 0x05, $trig, 0)
  $s = 0x48; for ($i = 0; $i -lt 4; $i++) {{ $s = ($s + $pl[$i]) -band 0xFF }}
  $pl[4] = [byte]$s
  $p.DiscardInBuffer(); $p.Write($pl, 0, 5)
  $ok = $false; $sw2 = [Diagnostics.Stopwatch]::StartNew()
  while ($sw2.ElapsedMilliseconds -lt 6000) {{ try {{ if ($p.ReadByte() -eq 6) {{ $ok = $true; break }} }} catch {{}} }}
  if (-not $ok) {{ Write-Output "ERR: ошибка стирания"; $p.Close(); exit 1 }}
}}
Write-Output "INFO: область приложения стёрта"

# --- дополнение до кратности 1024 ---
$pad = (1024 - ($fw.Length % 1024)) % 1024
if ($pad -gt 0) {{
  $t = New-Object byte[] ($fw.Length + $pad)
  [Array]::Copy($fw, $t, $fw.Length)
  for ($i = $fw.Length; $i -lt $t.Length; $i++) {{ $t[$i] = 0xFF }}
  $fw = $t
}}
$total = $fw.Length
$blocks = $total / 1024

# --- запись ---
$pl = New-Object byte[] 1028
for ($off = 0; $off -lt $total; $off += 1024) {{
  $pl[0] = 0x57
  $pl[1] = [byte](($off -shr 8) -band 0xFF)
  $pl[2] = [byte]($off -band 0xFF)
  [Array]::Copy($fw, $off, $pl, 3, 1024)
  $s = 0x48
  for ($i = 0; $i -lt 1027; $i++) {{ $s = ($s + $pl[$i]) -band 0xFF }}
  $pl[1027] = [byte]$s
  $p.DiscardInBuffer(); $p.Write($pl, 0, 1028)
  $ok = $false; $sw3 = [Diagnostics.Stopwatch]::StartNew()
  while ($sw3.ElapsedMilliseconds -lt 4000) {{ try {{ if ($p.ReadByte() -eq 6) {{ $ok = $true; break }} }} catch {{}} }}
  if (-not $ok) {{ Write-Output ("ERR: нет ACK на блоке 0x{{0:X5}}" -f $off); $p.Close(); exit 1 }}
  $n = ($off / 1024) + 1
  if (($n % 16) -eq 0 -or $n -eq $blocks) {{ Write-Output ("PROG: {{0}}/{{1}} блоков" -f $n, $blocks) }}
}}
$p.Close()
Write-Output "OK: записано. Выключите и включите рацию обычным образом."
'''


def find_port() -> str:
    out = subprocess.run(["powershell.exe", "-NoProfile", "-Command",
                          "[System.IO.Ports.SerialPort]::GetPortNames() -join ','"],
                         capture_output=True, text=True).stdout
    ports = [p for p in out.strip().replace("\r", "").split(",") if p]
    if not ports:
        sys.exit("COM-порт не найден — подключите рацию.")
    return ports[0]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("firmware")
    ap.add_argument("--port")
    a = ap.parse_args()

    fw = os.path.abspath(a.firmware)
    if not os.path.exists(fw):
        sys.exit(f"нет файла: {fw}")
    name = "rt4d_flash_tmp.bin"
    shutil.copy(fw, os.path.join(WIN_TMP_WSL, name))
    win_fw = os.path.join(WIN_TMP_WIN, name)

    port = a.port or find_port()
    script = PS.format(fwpath=win_fw.replace("\\", "\\\\"), port=port)

    proc = subprocess.Popen(["powershell.exe", "-NoProfile", "-Command", script],
                            stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                            text=True, encoding='utf-8', errors='replace')
    rc = 0
    for line in proc.stdout:
        line = line.rstrip("\r\n")
        if line:
            print(line, flush=True)
            if line.startswith("ERR:"):
                rc = 1
    proc.wait()
    sys.exit(rc or proc.returncode)


if __name__ == "__main__":
    main()