#!/usr/bin/env python3
"""
Simuliert einen Pico, der den OTA-Server anfragt. Damit testen wir das
Protokoll ohne echtes Geraet.

Usage:
  test_client.py <pico_id> [--current-binary PATH] [--host HOST] [--port PORT]

Beispiele:
  # so tun, als laege auf dem Pico ein leerer Slot (alles 0xFF):
  test_client.py 1

  # so tun, als laege auf dem Pico schon pico_1.bin:
  test_client.py 1 --current-binary firmware/pico_1.bin

  # unbekannte Pico-ID:
  test_client.py 99
"""

import argparse
import hashlib
import socket
import struct
import sys
from pathlib import Path

MAGIC = 0xC0FFEE01
APP_SLOT_SIZE = 3584 * 1024
SHA256_SIZE = 32


def slot_hash(binary_path: Path | None) -> bytes:
    if binary_path is None:
        # Leerer Slot: 3.5 MB 0xFF.
        padded = b"\xff" * APP_SLOT_SIZE
    else:
        data = binary_path.read_bytes()
        if len(data) > APP_SLOT_SIZE:
            raise ValueError(f"{binary_path}: zu gross ({len(data)})")
        padded = data + b"\xff" * (APP_SLOT_SIZE - len(data))
    return hashlib.sha256(padded).digest()


def recv_exact(sock: socket.socket, n: int) -> bytes:
    buf = b""
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            raise ConnectionError(f"Server-Disconnect nach {len(buf)}/{n}")
        buf += chunk
    return buf


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("pico_id", type=int)
    ap.add_argument("--current-binary", type=Path, default=None,
                    help="Pfad einer Binary, die wir 'als bereits installiert' simulieren")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, default=9000)
    ap.add_argument("--save-update", type=Path, default=None,
                    help="Falls Server ein Update sendet, hier speichern")
    args = ap.parse_args()

    current_hash = slot_hash(args.current_binary)
    print(f"-> pico_id={args.pico_id}  hash={current_hash.hex()[:16]}...")

    req = struct.pack("<II", MAGIC, args.pico_id) + current_hash
    assert len(req) == 40

    with socket.create_connection((args.host, args.port), timeout=5.0) as sock:
        sock.sendall(req)

        status_byte = recv_exact(sock, 1)
        status = status_byte[0]
        print(f"<- status=0x{status:02x}")

        if status == 0x00:
            print("   no update — boot the existing app")
            return 0

        if status != 0x01:
            print(f"   FEHLER: unerwartetes Status-Byte 0x{status:02x}")
            return 1

        size_bytes = recv_exact(sock, 4)
        size = struct.unpack("<I", size_bytes)[0]
        new_hash = recv_exact(sock, SHA256_SIZE)
        print(f"   size={size}  new_hash={new_hash.hex()[:16]}...")

        payload = recv_exact(sock, size)
        actual_slot_hash = hashlib.sha256(
            payload + b"\xff" * (APP_SLOT_SIZE - len(payload))
        ).digest()
        match = actual_slot_hash == new_hash
        print(f"   payload empfangen ({len(payload)} bytes)  "
              f"hash-match={match}")

        if not match:
            print("   FEHLER: gepaddeter Payload-Hash passt nicht zum angekuendigten Slot-Hash")
            return 2

        if args.save_update:
            args.save_update.write_bytes(payload)
            print(f"   gespeichert: {args.save_update}")

    return 0


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