#!/usr/bin/env python3
"""
hex_source.py — holt den eingebetteten MakeCode-Quelltext aus einer micro:bit .hex.

Jede von MakeCode erzeugte .hex trägt das komplette Projekt (main.ts, alle
Extensions, pxt.json) LZMA-komprimiert am Ende mit. Genau daraus stammen die
echten Kosmos-Proxi-Quellen inklusive der Hardware-Extension `proxi.ts` — die
Pin-Belegung des Roboters ist sonst nirgends dokumentiert (siehe
docs/HARDWARE.md).

Funktioniert mit jeder MakeCode-.hex, nicht nur mit denen von Proxi.

    python3 hex_source.py *.hex -o ./quellen
    python3 hex_source.py --list 00_Funktionstest.hex
    python3 hex_source.py --check firmware.hex main.ts

Format des Anhangs (ab der Magic-Sequenz):
    +0   8 Byte   Magic 41 14 0E 2F B8 2F A2 BB
    +8   2 Byte   Länge des JSON-Kopfes (little endian)
    +10  4 Byte   Länge der komprimierten Nutzdaten (little endian)
    +14  2 Byte   reserviert
    +16  n Byte   JSON-Kopf {"compression":"LZMA","textSize":...,"name":...}
    +16+n         LZMA-Strom (alone/legacy-Header) ->
                  JSON-Kopf (Projekt-Metadaten) + JSON-Objekt {dateiname: inhalt}
"""

import argparse
import binascii
import json
import lzma
import os
import sys

MAGIC = bytes.fromhex("41140E2FB82FA2BB")

# Intel-HEX-Record-Typen, die Nutzdaten tragen. 0x0D/0x0E sind Erweiterungen des
# micro:bit-"Universal Hex" (eine Datei für V1 und V2), 0x0C ist reine Füllung.
DATA_RECORDS = {0x00, 0x0C, 0x0D, 0x0E}


def _hex_payload(path):
    """Alle Datenbytes der .hex hintereinander — ohne Adressauswertung.

    Reicht hier, weil der Quelltext-Anhang am Stück in aufeinanderfolgenden
    Records liegt und ab der Magic-Sequenz gelesen wird.
    """
    out = bytearray()
    with open(path, "r", errors="replace") as fh:
        for line in fh:
            line = line.strip()
            if not line.startswith(":"):
                continue
            try:
                raw = binascii.unhexlify(line[1:])
            except binascii.Error:
                continue
            if len(raw) < 5:
                continue
            length, rectype = raw[0], raw[3]
            if rectype in DATA_RECORDS:
                out += raw[4:4 + length]
    return bytes(out)


def read_project(path):
    """(metadaten, {dateiname: inhalt}) aus einer .hex lesen."""
    blob = _hex_payload(path)
    start = blob.find(MAGIC)
    if start < 0:
        raise ValueError("kein eingebetteter MakeCode-Quelltext gefunden")

    header_len = int.from_bytes(blob[start + 8:start + 10], "little")
    packed_len = int.from_bytes(blob[start + 10:start + 14], "little")
    packed_at = start + 16 + header_len

    text = lzma.decompress(
        blob[packed_at:packed_at + packed_len], format=lzma.FORMAT_ALONE
    ).decode("utf-8", "replace")

    decoder = json.JSONDecoder()
    meta, offset = decoder.raw_decode(text)
    files, _ = decoder.raw_decode(text[offset:].lstrip())
    return meta, files


def cmd_extract(paths, outdir):
    for path in paths:
        name = os.path.splitext(os.path.basename(path))[0]
        try:
            meta, files = read_project(path)
        except Exception as exc:
            print(f"!! {name}: {exc}", file=sys.stderr)
            continue
        target = os.path.join(outdir, name)
        os.makedirs(target, exist_ok=True)
        for filename, content in files.items():
            dest = os.path.join(target, filename)
            os.makedirs(os.path.dirname(dest), exist_ok=True)
            with open(dest, "w") as fh:
                fh.write(content)
        print(f"{name}: {', '.join(sorted(files))}  ->  {target}")


def cmd_list(paths):
    for path in paths:
        meta, files = read_project(path)
        print(f"== {os.path.basename(path)}  ({meta.get('name', '?')})")
        for filename in sorted(files):
            print(f"   {filename:<16} {len(files[filename]):>7} Zeichen")


def _normalize(text):
    """Zeilenenden vereinheitlichen.

    MakeCode bettet den Quelltext so ein, wie er auf der Platte lag — bei
    Dateien mit CRLF also mit CRLF. Beim Lesen übersetzt Python das
    stillschweigend nach LF. Ohne Normalisierung meldet der Vergleich deshalb
    einen Unterschied, den es gar nicht gibt.
    """
    return text.replace("\r\n", "\n").replace("\r", "\n").strip()


def cmd_check(hexfile, source):
    """Prüft, ob die .hex genau den übergebenen Quelltext enthält."""
    _, files = read_project(hexfile)
    name = os.path.basename(source)
    if name not in files:
        print(f"FEHLER: {name} steckt nicht in {hexfile}", file=sys.stderr)
        return 1
    with open(source) as fh:
        expected = _normalize(fh.read())
    found = _normalize(files[name])
    if found != expected:
        print(f"FEHLER: {name} in {hexfile} weicht von {source} ab "
              f"({len(found)} vs. {len(expected)} Zeichen) — "
              f".hex ist veraltet.", file=sys.stderr)
        return 1
    print(f"OK: {hexfile} enthält den aktuellen {name}")
    return 0


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("files", nargs="+", help=".hex-Dateien (bei --check: <hex> <quelldatei>)")
    ap.add_argument("-o", "--outdir", default=".", help="Zielordner für die Extraktion")
    ap.add_argument("--list", action="store_true", help="nur auflisten, nichts schreiben")
    ap.add_argument("--check", action="store_true",
                    help="prüfen, ob <hex> den aktuellen <quelldatei>-Stand enthält")
    args = ap.parse_args()

    if args.check:
        if len(args.files) != 2:
            ap.error("--check braucht genau zwei Argumente: <hex> <quelldatei>")
        return cmd_check(args.files[0], args.files[1])
    if args.list:
        cmd_list(args.files)
        return 0
    cmd_extract(args.files, args.outdir)
    return 0


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