#!/usr/bin/env python3
"""Emit a spec-correct 32-bit LE ELF (e_machine=EM_NDS32) from a MediaTek
Connac2 WiFi firmware container, one PT_LOAD segment per region at true
load addresses.

Structure is validated by readelf; disassembly requires Ghidra (official
NDS32 module) or an Andes nds32le-elf toolchain — Debian binutils dropped
NDS32. Entry point is unknown pre-RE (e_entry=0); set it in your tool after
import.

NON_DL regions are included as loadable segments by default; they are
host-side data per the driver, not downloaded to the MCU.
"""
import struct
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from mtk_fw_extract import PATCH_HDR, PATCH_SEC, REGION, TRAILER  # noqa: E402

EM_NDS32 = 167
EHDR = struct.Struct('<16sHHIIIIIHHHHHH')  # 52 bytes
PHDR = struct.Struct('<IIIIIIII')          # 32 bytes
PT_LOAD, PF_RWX = 1, 7


def ram_regions(d: bytes):
    """Yield (addr, data) for each region of a RAM container."""
    t = TRAILER.unpack_from(d, len(d) - TRAILER.size)
    n = t[2]
    table = len(d) - TRAILER.size - n * REGION.size
    off = 0
    for i in range(n):
        r = REGION.unpack_from(d, table + i * REGION.size)
        addr, ln = r[4], r[5]
        data = d[off:off + ln]
        off += ln
        yield addr, data


def patch_regions(d: bytes):
    """Yield (addr, data) for each section of a patch container."""
    h = PATCH_HDR.unpack_from(d, 0)
    n = h[9]
    for i in range(n):
        s = PATCH_SEC.unpack_from(d, PATCH_HDR.size + i * PATCH_SEC.size)
        typ, offs, size = s[0], s[1], s[2]
        addr = s[3]
        yield addr, d[offs:offs + size]


def build_elf(regions) -> bytes:
    regions = [(a, dat) for a, dat in regions if dat]
    n = len(regions)
    phoff = EHDR.size
    data_off = phoff + n * PHDR.size
    eh = EHDR.pack(
        b'\x7fELF' + bytes([1, 1, 1, 0]) + b'\x00' * 8,  # ELFCLASS32/LSB/v1
        2, EM_NDS32, 1,
        regions[0][0] if regions else 0,   # e_entry (first region; see doc)
        phoff, 0, 0,
        EHDR.size, PHDR.size, n,
        40, 0, 0)
    out = bytearray(eh)
    offset = data_off
    for addr, dat in regions:
        out += PHDR.pack(PT_LOAD, offset, addr, addr, len(dat), len(dat),
                         PF_RWX, 0x1000)
        offset += len(dat)
    for _, dat in regions:
        out += dat
    return bytes(out)


def main():
    if len(sys.argv) < 2:
        sys.exit(f'usage: {sys.argv[0]} <blob> [<blob>...] [-o outdir]')
    args, outdir = [], Path('extracted')
    it = iter(sys.argv[1:])
    for a in it:
        if a == '-o':
            outdir = Path(next(it))
        else:
            args.append(a)
    for path in args:
        p = Path(path)
        d = p.read_bytes()
        regions = patch_regions(d) if 'patch' in p.name else ram_regions(d)
        elf = build_elf(regions)
        out = outdir / p.stem / f'{p.stem}.elf'
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_bytes(elf)
        n_seg = struct.unpack_from('<H', elf, 44)[0]  # e_phnum
        print(f'{p.name} -> {out} ({len(elf)} bytes, {n_seg} PT_LOAD)')


if __name__ == '__main__':
    main()