#!/usr/bin/env python3
"""
OTA-Update-Server fuer den Pico-Bootloader.

Wire-Protokoll:
  Pico -> Server (40 bytes):
    [0:4]   magic 0xC0FFEE01 (LE)
    [4:8]   pico_id (LE u32)
    [8:40]  SHA256 des aktuellen App-Slots (3.5 MB, 0xFF-gepaddet)

  Server -> Pico:
    [0:1]   status: 0x00 = no update, 0x01 = update follows
    falls status == 0x01:
      [1:5]   size der Binary in bytes (LE u32)
      [5:37]  SHA256 des neuen App-Slots (auch 0xFF-gepaddet auf 3.5 MB)
      [37:]   Binary-Payload (size Bytes)
"""

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

MAGIC = 0xC0FFEE01
APP_SLOT_SIZE = 3584 * 1024  # 3.5 MB
SHA256_SIZE = 32
REQUEST_SIZE = 4 + 4 + SHA256_SIZE  # 40

STATUS_NO_UPDATE = 0x00
STATUS_UPDATE = 0x01

LISTEN_HOST = "0.0.0.0"
LISTEN_PORT = 9000
RECV_TIMEOUT_S = 5.0

FIRMWARE_DIR = Path(__file__).parent / "firmware"

# pico_id -> binary path. Picos, die hier nicht gelistet sind, bekommen
# immer status=no_update (server kennt das Geraet nicht).
MANIFEST = {
    1: FIRMWARE_DIR / "pico_2.bin",
    2: FIRMWARE_DIR / "pico_1.bin",
}


def slot_hash(binary_path: Path) -> bytes:
    """Hash der Datei, gepaddet mit 0xFF bis 3.5 MB — exakt das,
    was der Pico auf seinem App-Slot berechnet."""
    data = binary_path.read_bytes()
    if len(data) > APP_SLOT_SIZE:
        raise ValueError(
            f"{binary_path}: {len(data)} bytes > app slot ({APP_SLOT_SIZE})"
        )
    padded = data + b"\xff" * (APP_SLOT_SIZE - len(data))
    return hashlib.sha256(padded).digest()


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


def handle_client(conn: socket.socket, addr) -> None:
    conn.settimeout(RECV_TIMEOUT_S)

    try:
        request = recv_exact(conn, REQUEST_SIZE)
    except (ConnectionError, socket.timeout) as exc:
        print(f"[{addr}] Request-Abbruch: {exc}")
        return

    magic, pico_id = struct.unpack("<II", request[:8])
    current_hash = request[8:]

    print(
        f"[{addr}] pico_id={pico_id} hash={current_hash.hex()[:16]}... magic=0x{magic:08x}"
    )

    if magic != MAGIC:
        print(f"[{addr}] FEHLER: ungueltiges Magic")
        return

    binary_path = MANIFEST.get(pico_id)
    if binary_path is None:
        print(f"[{addr}] pico_id {pico_id} nicht im Manifest -> no_update")
        conn.sendall(bytes([STATUS_NO_UPDATE]))
        return

    if not binary_path.exists():
        print(f"[{addr}] Manifest-Eintrag fehlt: {binary_path} -> no_update")
        conn.sendall(bytes([STATUS_NO_UPDATE]))
        return

    target_hash = slot_hash(binary_path)

    if current_hash == target_hash:
        print(f"[{addr}] aktuell -> no_update")
        conn.sendall(bytes([STATUS_NO_UPDATE]))
        return

    binary = binary_path.read_bytes()
    print(
        f"[{addr}] update -> {binary_path.name} ({len(binary)} bytes, "
        f"hash {target_hash.hex()[:16]}...)"
    )
    conn.sendall(bytes([STATUS_UPDATE]))
    conn.sendall(struct.pack("<I", len(binary)))
    conn.sendall(target_hash)
    conn.sendall(binary)


def main() -> int:
    if not FIRMWARE_DIR.exists():
        FIRMWARE_DIR.mkdir(parents=True)
        print(f"firmware/ angelegt: {FIRMWARE_DIR}")

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as srv:
        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        srv.bind((LISTEN_HOST, LISTEN_PORT))
        srv.listen(4)
        print(f"OTA-Server hoert auf {LISTEN_HOST}:{LISTEN_PORT}")
        print(f"Manifest: {[(k, str(v.name)) for k, v in MANIFEST.items()]}")
        try:
            while True:
                conn, addr = srv.accept()
                try:
                    handle_client(conn, addr)
                except Exception as exc:
                    print(f"[{addr}] Exception: {exc}")
                finally:
                    conn.close()
        except KeyboardInterrupt:
            print("\nbeendet")
            return 0


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