"""
usb_esptool_push.py — cliente USB v2 do Albert (substitui usb_k2_push.py).

Substitui o protocolo USB-K2 custom por: esptool stub oficial + plaintext via
HTTPS streaming. Plaintext NUNCA toca disco — fica em bytearray na RAM Python,
chunks zerados conforme enviados. K3 hardware (eFuse BLK1) cifra ao gravar.

Fluxo:
  1. Serial app (115200): manda "whoami\\n", parseia "running=app0/app1".
  2. Fecha serial.
  3. HTTP GET /usb_update_bundle/<NS>     → target_version, plaintext_url, meta_url.
  4. HTTP GET <meta_url>                  → {sha256_plain, size_plain}.
  5. HTTP GET <plaintext_url> streaming   → bytearray em RAM; sha256 incremental.
  6. esptool.detect_chip(port)            → ESP32 em download mode via DTR/RTS.
  7. esp.run_stub()                       → IS_STUB=True (encrypted_write requer stub).
  8. esptool.attach_flash(esp)            → identifica SPI flash.
  9. esp.flash_begin(size, offset, encrypted_write=True).
 10. Loop FLASH_WRITE_SIZE chunks: flash_block(chunk, seq, encrypted=True);
     zera chunk no buffer apos enviar.
 11. esp.flash_finish(reboot=False).
 12. esp.hard_reset().
"""
import argparse
import gc
import hashlib
import json
import sys
import time
import urllib.request
from pathlib import Path

try:
    import serial
except ImportError:
    print("[ERRO] pyserial nao instalado. Rode: pip install pyserial", file=sys.stderr)
    sys.exit(2)

try:
    import esptool
except ImportError:
    print("[ERRO] esptool lib nao instalada. Rode: pip install esptool", file=sys.stderr)
    sys.exit(2)


# Offsets das particoes Arduino default (Minimal SPIFFS Large APPS with OTA):
# Albert usa esse layout em todas as v1.8.x. Confirmado em
# project_albert_build_conventions.md.
PARTITION_OFFSETS = {
    "app0": 0x10000,
    "app1": 0x1F0000,
}

DEFAULT_BAUD_SERIAL = 115200    # app mode
DEFAULT_BAUD_CONNECT = 115200   # SYNC inicial com ROM bootloader (sempre 115200)
DEFAULT_BAUD_FLASH  = 460800    # change_baud depois do stub (~50 KB/s real)
DEFAULT_BASE_URL    = "http://iot.mgconnect.online"

WHOAMI_TIMEOUT_S    = 10.0
HTTP_TIMEOUT_S      = 60.0


def parse_whoami_line(text: str) -> dict:
    """'WHOAMI: numSerie=X mod=Y running=app0 ...' -> dict."""
    if ":" not in text:
        return {}
    out = {}
    body = text.split(":", 1)[1].strip()
    for token in body.split():
        if "=" in token:
            k, v = token.split("=", 1)
            out[k] = v
    return out


def force_app_boot(port: str) -> None:
    """Pulso RTS pra forcar reset do ESP32 com GPIO0=HIGH (boot em app mode).

    Usado quando whoami nao responde — Albert pode estar parado num estado
    quirky apos um esptool falho na rodada anterior. Reset coloca app pra
    rodar de novo.
    """
    s = serial.Serial()
    s.port = port
    s.baudrate = DEFAULT_BAUD_SERIAL
    s.timeout = 0.1
    s.dtr = False
    s.rts = False
    s.open()
    try:
        s.dtr = False   # GPIO0 livre (pull-up = HIGH = app mode)
        s.rts = True    # EN LOW = reset
        time.sleep(0.2)
        s.rts = False   # EN HIGH = sai do reset, boota app
    finally:
        try:
            s.close()
        except Exception:
            pass
    time.sleep(2.0)  # Espera boot do app completar antes do whoami


def read_whoami(port: str) -> dict:
    """Abre serial app, manda whoami, lê primeira linha WHOAMI:. Fecha serial."""
    s = serial.Serial()
    s.port = port
    s.baudrate = DEFAULT_BAUD_SERIAL
    s.timeout = 0.5
    s.dtr = False
    s.rts = False
    s.open()
    try:
        s.reset_input_buffer()
        s.write(b"whoami\n")
        s.flush()
        deadline = time.time() + WHOAMI_TIMEOUT_S
        buf = b""
        while time.time() < deadline:
            chunk = s.read(256)
            if not chunk:
                continue
            buf += chunk
            while b"\n" in buf:
                idx = buf.index(b"\n")
                line = buf[:idx].decode("utf-8", errors="replace").rstrip("\r")
                buf = buf[idx + 1:]
                if line.startswith("WHOAMI:"):
                    return parse_whoami_line(line)
        return {}
    finally:
        try:
            s.close()
        except Exception:
            pass


def http_get_json(url: str) -> dict:
    with urllib.request.urlopen(url, timeout=HTTP_TIMEOUT_S) as resp:
        return json.loads(resp.read())


def http_stream_to_ram(url: str, expected_size: int, expected_sha256_hex: str) -> bytearray:
    """Baixa firmware plaintext via HTTPS streaming pra bytearray em RAM.

    Valida sha256 incremental + Content-Length + X-Albert-SHA256 do servidor.
    Aborta antes de tocar no device se algo nao bater.
    """
    print(f"[push] HTTP GET {url}")
    sha = hashlib.sha256()
    buf = bytearray()

    with urllib.request.urlopen(url, timeout=HTTP_TIMEOUT_S) as resp:
        cl = resp.headers.get("Content-Length")
        if cl:
            cl_int = int(cl)
            if cl_int != expected_size:
                raise RuntimeError(
                    f"Content-Length={cl_int} != esperado {expected_size}"
                )

        srv_sha = (resp.headers.get("X-Albert-SHA256", "") or "").lower()
        if srv_sha and srv_sha != expected_sha256_hex.lower():
            raise RuntimeError(
                f"X-Albert-SHA256={srv_sha} != esperado {expected_sha256_hex}"
            )

        t0 = time.time()
        last_ts = t0
        last_bytes = 0
        while True:
            chunk = resp.read(65536)
            if not chunk:
                break
            buf.extend(chunk)
            sha.update(chunk)
            now = time.time()
            if now - last_ts >= 1.0:
                rate = (len(buf) - last_bytes) / (now - last_ts)
                pct = 100.0 * len(buf) / max(expected_size, 1)
                print(
                    f"  [push] HTTP {len(buf)}/{expected_size} bytes "
                    f"({pct:.1f}%) — {rate/1024:.1f} KB/s",
                    flush=True,
                )
                last_ts = now
                last_bytes = len(buf)

    if len(buf) != expected_size:
        raise RuntimeError(f"size={len(buf)} != esperado {expected_size}")

    got_sha = sha.hexdigest().lower()
    if got_sha != expected_sha256_hex.lower():
        raise RuntimeError(f"sha256={got_sha} != esperado {expected_sha256_hex}")

    print(f"[push] HTTP OK — {len(buf)} bytes, sha256 validado")
    return buf


def flash_via_esptool(port: str, plaintext: bytearray, offset: int) -> None:
    """Conecta via esptool, sobe stub, grava com encrypted_write.

    Após cada chunk enviado, zera o intervalo correspondente em plaintext.
    O buffer nunca chega ao disco; ficam em RAM Python até ser zerado.
    """
    # ROM bootloader SYNC sempre em 115200. Conectar em baud alto causa "No serial
    # data received" porque o handshake inicial nao consegue completar antes do
    # ROM dar timeout. So aceleramos depois do stub via change_baud.
    # connect_attempts=10 (default 7) cobre jitter no reset DTR/RTS.
    print(f"[push] esptool: detect_chip {port} @ {DEFAULT_BAUD_CONNECT} baud")
    esp = esptool.detect_chip(port=port, baud=DEFAULT_BAUD_CONNECT, connect_attempts=10)
    print(f"[push] esptool: chip={esp.CHIP_NAME}")

    print("[push] esptool: run_stub")
    esp = esp.run_stub()

    if DEFAULT_BAUD_FLASH != DEFAULT_BAUD_CONNECT:
        print(f"[push] esptool: change_baud {DEFAULT_BAUD_FLASH}")
        esp.change_baud(DEFAULT_BAUD_FLASH)

    print("[push] esptool: attach_flash")
    esptool.attach_flash(esp)

    block_size = esp.FLASH_WRITE_SIZE
    align = esp.FLASH_ENCRYPTED_WRITE_ALIGN
    size_real = len(plaintext)
    # Padding pra alinhar ao block_size (resto vira 0xFF).
    # ESP32 boot ROM le size do header do app — bytes extras 0xFF nao atrapalham.
    size_padded = ((size_real + block_size - 1) // block_size) * block_size

    print(
        f"[push] esptool: flash_begin offset=0x{offset:X} size={size_padded} "
        f"(real={size_real}, block={block_size}, align={align})"
    )
    num_blocks = esp.flash_begin(size_padded, offset, encrypted_write=True)
    print(f"[push] esptool: erase OK, num_blocks={num_blocks}")

    t0 = time.time()
    last_ts = t0
    sent = 0
    seq = 0
    while sent < size_padded:
        # Monta chunk de exatamente block_size bytes (pad com 0xFF no fim).
        end_real = min(sent + block_size, size_real)
        if end_real > sent:
            chunk = bytes(plaintext[sent:end_real])
            real_bytes = end_real - sent
        else:
            chunk = b""
            real_bytes = 0
        if real_bytes < block_size:
            chunk = chunk + b"\xFF" * (block_size - real_bytes)

        esp.flash_block(chunk, seq, encrypted=True)

        # Zera os bytes reais recem-enviados (0xFF de padding sao constante).
        for i in range(sent, end_real):
            plaintext[i] = 0

        sent += block_size
        seq += 1
        now = time.time()
        if now - last_ts >= 1.0 or sent >= size_padded:
            elapsed = now - t0
            rate = sent / max(elapsed, 0.001)
            pct = 100.0 * sent / size_padded
            print(
                f"  [push] flash {sent}/{size_padded} bytes "
                f"({pct:.1f}%) — {rate/1024:.1f} KB/s",
                flush=True,
            )
            last_ts = now

    esp.flash_finish(reboot=False)
    elapsed = time.time() - t0
    print(
        f"[push] esptool: flash OK em {elapsed:.1f}s "
        f"({size_padded/elapsed/1024:.1f} KB/s)"
    )

    print("[push] esptool: hard_reset (DTR/RTS)")
    esp.hard_reset()


def main():
    ap = argparse.ArgumentParser(
        description="Push firmware to Albert via esptool stub + plaintext HTTPS"
    )
    ap.add_argument("--port", required=True, help="Porta COM (ex: COM4)")
    ap.add_argument("--numserie", required=True, help="Numero de serie do device (ex: 20LBE)")
    ap.add_argument("--base-url", default=DEFAULT_BASE_URL,
                    help=f"Base URL do AlbertService (default {DEFAULT_BASE_URL})")
    ap.add_argument("--bootstrap", action="store_true",
                    help="Modo bootstrap: ESP32 com K3 burnado mas SEM firmware Albert. "
                         "Pula whoami, grava em offset --offset (default app0=0x10000). "
                         "Depois do reset, rode 2_Primeira_Instalacao pra popular NVS.")
    ap.add_argument("--offset", type=lambda s: int(s, 0), default=0x10000,
                    help="Offset de gravacao em modo --bootstrap (default 0x10000=app0). "
                         "Ignorado em modo normal (offset vem do whoami running=).")
    args = ap.parse_args()

    # ETAPA 1: descobrir offset
    if args.bootstrap:
        # Bootstrap: ESP32 com K3 ativo mas sem Albert. Sem whoami.
        offset = args.offset
        print(f"[push] modo BOOTSTRAP: numSerie={args.numserie}, "
              f"offset=0x{offset:X} (sem whoami — Albert nao roda ainda)")
    else:
        # Normal: whoami pra descobrir running partition.
        print(f"[push] whoami em {args.port}...")
        info = read_whoami(args.port)
        if not info:
            # Auto-recuperacao: Albert pode estar parado apos um esptool falho
            # anterior. Pulso RTS pra forcar reset em app mode + retry.
            print("[push] whoami sem resposta, tentando reset RTS + retry...")
            try:
                force_app_boot(args.port)
            except Exception as e:
                print(f"[ERRO] reset falhou: {e}", file=sys.stderr)
                sys.exit(1)
            info = read_whoami(args.port)
            if not info:
                print("[ERRO] whoami nao respondeu mesmo apos reset. Outra janela "
                      "na COM? Cabo USB? Firmware travado? Se ESP32 nunca teve "
                      "Albert, use --bootstrap.", file=sys.stderr)
                sys.exit(1)

        ns_dev = info.get("numSerie", "")
        if ns_dev != args.numserie:
            print(f"[ERRO] numSerie do device ({ns_dev!r}) != --numserie ({args.numserie!r})",
                  file=sys.stderr)
            sys.exit(2)

        running = info.get("running", "")
        if running not in PARTITION_OFFSETS:
            print(f"[ERRO] whoami running={running!r} (esperado app0/app1). "
                  f"Firmware tem o campo running no whoami?", file=sys.stderr)
            sys.exit(3)

        offset = PARTITION_OFFSETS[running]
        sw_atual = info.get("sw", "")
        print(f"[push] device {info.get('numSerie','')}: rodando {running} "
              f"(offset 0x{offset:X}), SW={sw_atual}")

    # ETAPA 2: bundle pra descobrir target + plaintext_url + meta_url
    bundle_url = f"{args.base_url}/service/api/usb_update_bundle/{args.numserie}"
    print(f"[push] HTTP GET {bundle_url}")
    try:
        bundle = http_get_json(bundle_url)
    except Exception as e:
        print(f"[ERRO] /usb_update_bundle falhou: {e}", file=sys.stderr)
        sys.exit(4)

    if not bundle.get("ok"):
        print(f"[ERRO] bundle invalido: {bundle}", file=sys.stderr)
        sys.exit(4)

    target = bundle.get("target_version", "")
    plaintext_url = bundle.get("plaintext_url", "")
    meta_url = bundle.get("meta_url", "")
    if not target or not plaintext_url or not meta_url:
        print(f"[ERRO] bundle faltando campos: target={target!r} "
              f"plaintext_url={plaintext_url!r} meta_url={meta_url!r}",
              file=sys.stderr)
        sys.exit(4)

    print(f"[push] target version: {target}")

    # ETAPA 3: meta → sha256_plain + size_plain
    print(f"[push] HTTP GET {meta_url}")
    try:
        meta = http_get_json(meta_url)
    except Exception as e:
        print(f"[ERRO] /ota_metadata falhou: {e}", file=sys.stderr)
        sys.exit(4)

    expected_sha = (meta.get("sha256_plain") or "").lower()
    expected_size = int(meta.get("size_plain") or 0)
    if not expected_sha or expected_size <= 0:
        print(f"[ERRO] meta invalido: sha256_plain={expected_sha!r} "
              f"size_plain={expected_size}", file=sys.stderr)
        sys.exit(4)

    print(f"[push] esperado: {expected_size} bytes, sha256={expected_sha[:8]}...")

    # ETAPA 4: download plaintext streaming → RAM
    try:
        plaintext = http_stream_to_ram(plaintext_url, expected_size, expected_sha)
    except Exception as e:
        print(f"[ERRO] download plaintext falhou: {e}", file=sys.stderr)
        sys.exit(5)

    # ETAPA 5: esptool flash com encrypted_write
    # Em modo normal: delay pra ESP32 estabilizar apos o close da serial do
    # whoami — abrir/fechar COM no Windows pulsa DTR/RTS, e o esptool precisa
    # do device em estado previsivel pra o reset funcionar.
    # Em modo bootstrap nao houve whoami, mas mantemos o delay (defensivo: a
    # porta foi aberta brevemente no force_app_boot ou em outra logica? Nao;
    # bootstrap pula tudo isso, mas custo de 1s e irrelevante).
    if not args.bootstrap:
        time.sleep(1.0)
    try:
        try:
            flash_via_esptool(args.port, plaintext, offset)
        finally:
            # Zera qualquer byte que tenha sobrado (defensivo).
            for i in range(len(plaintext)):
                plaintext[i] = 0
            del plaintext
            gc.collect()
    except Exception as e:
        print(f"[ERRO] esptool flash falhou: {e}", file=sys.stderr)
        sys.exit(6)

    print(f"[push] OK — Albert reiniciando com firmware novo "
          f"em {time.strftime('%H:%M:%S')}")
    return 0


if __name__ == "__main__":
    sys.exit(main())
