#!/usr/bin/env python3
"""Pushes a signed Update File to a Cardputer over Wi-Fi (TCP 3232) and reports the result.

Usage: scripts/ota_push.py <file.ota> <host>
<host> is the device's IP (shown in Settings > Firmware).
"""
import socket
import sys

PORT = 3232


def main():
    if len(sys.argv) < 3:
        sys.exit(__doc__)
    path, host = sys.argv[1:3]
    data = open(path, "rb").read()
    try:
        reply = push(data, host)
    except (ConnectionResetError, BrokenPipeError):
        # The device checks the header first and hangs up on a refused update mid-transfer.
        print()
        print("device: refused the update and closed the connection (the reason is on its screen)")
        sys.exit(1)
    print(f"device: {reply or '(no answer)'}")
    sys.exit(0 if reply.startswith("OK") else 1)


def push(data, host):
    with socket.create_connection((host, PORT), timeout=15) as s:
        s.settimeout(60)
        sent, last = 0, -1
        while sent < len(data):
            chunk = data[sent:sent + 4096]
            s.sendall(chunk)
            sent += len(chunk)
            pct = sent * 100 // len(data)
            if pct != last:
                print(f"\rsending {pct:3d}%", end="", flush=True)
                last = pct
        # The header announces the image size, so current firmware answers once it has it all.
        # Firmware from before that change only knows the file ended when we half-close.
        reply = b""
        s.settimeout(3)
        try:
            reply = s.recv(256)
        except socket.timeout:
            s.shutdown(socket.SHUT_WR)
            s.settimeout(60)
        while reply == b"" or not reply.endswith(b"\n"):
            try:
                part = s.recv(256)
            except socket.timeout:
                break
            if not part:
                break
            reply += part
    print()
    return reply.decode(errors="replace").strip()


if __name__ == "__main__":
    main()