"""Shared parser/packer for Lofi-12 XT Sonicware .bin images.

Format: [cmtd 48B][mmtd 36B][sect xN -> EOF, exact fit]
Each sect: b'sect' | u32 LE load_addr | u32 LE len | payload[len]

All addresses are DDR2 (0xC2xxxxxx, TMS320C6748). Code is C674x DSP LE.
Stdlib only.
"""
import struct
import json
import binascii

CMTD_LEN = 48
MMTD_LEN = 36

# Checksum seed: cmtd+0x08 coverage replaces file bytes [8:12] with this
# constant before CRC32-IEEE (init 0). Proven 3/3 against stock images
# (solved independently per image via GF(2), all give this same value;
# it is also built literally in the bootloader: MVK 0x6282/MVKH 0xC27C).
CKSEED = 0xC27C6282


def compute_checksum(d: bytes) -> int:
    """Valid cmtd+0x08 for a complete image: CRC32 of the image with
    bytes [8:12] replaced by CKSEED."""
    b = bytearray(d)
    b[8:12] = struct.pack('<I', CKSEED)
    return binascii.crc32(bytes(b)) & 0xFFFFFFFF

def parse_image(d: bytes) -> dict:
    assert d[0:4] == b'cmtd', "bad cmtd magic"
    assert d[0x30:0x34] == b'mmtd', "bad mmtd magic"
    fsize, cksum = struct.unpack('<II', d[4:12])
    assert fsize == len(d), f"cmtd filesize {fsize} != actual {len(d)}"
    cmtd = {
        'filesize': fsize,
        'checksum': cksum,
        'raw_0x0C_0x30': d[0x0C:0x30].hex(),
        'u32_0x10_0x30': list(struct.unpack('<8I', d[0x10:0x30])),
    }
    mmtd_rem, mmtd_flags, mmtd_ff = struct.unpack('<III', d[0x34:0x40])
    fw = struct.unpack('<III', d[0x40:0x4C])
    entry, nsect = struct.unpack('<II', d[0x4C:0x54])
    mmtd = {
        'remaining': mmtd_rem,
        'flags': mmtd_flags,
        'ff': mmtd_ff,
        'fw': list(fw),
        'entry': entry,
        'nsect': nsect,
    }
    assert mmtd_rem == len(d) - CMTD_LEN, "mmtd remaining mismatch"
    sects = []
    off = 0x54
    for i in range(nsect):
        assert d[off:off+4] == b'sect', f"bad sect magic at {off:#x}"
        addr, ln = struct.unpack('<II', d[off+4:off+12])
        payload = d[off+12:off+12+ln]
        assert len(payload) == ln, f"sect {i} truncated"
        sects.append({'index': i, 'off': off, 'addr': addr, 'len': ln,
                      'end': addr + ln, 'payload': payload})
        off += 12 + ln
    assert off == len(d), f"sect chain ends at {off:#x}, EOF {len(d):#x}"
    return {'cmtd': cmtd, 'mmtd': mmtd, 'sects': sects}


def build_image(meta: dict, payloads: list, checksum: int | None = None) -> bytes:
    """Rebuild image from meta (cmtd/mmtd dicts) + payload list [(addr, bytes)].
    Recalculates filesize/remaining fields. Preserves entry/flags/fw unless
    caller edited meta. Checksum: explicit value, 'keep' preserves the
    original from meta, None (default) computes the valid checksum."""
    nsect = len(payloads)
    total = CMTD_LEN + MMTD_LEN + sum(12 + len(p) for _, p in payloads)
    # checksum covers the whole file, so build with a zero placeholder,
    # then finalize: explicit int wins, 'keep' preserves meta, None computes.
    ck = 0
    out = bytearray()
    out += b'cmtd'
    out += struct.pack('<II', total, ck)
    out += bytes.fromhex(meta['cmtd']['raw_0x0C_0x30'])
    # fix embedded filesize/remaining inside cmtd raw (last two u32s):
    # raw layout: 8 x u32 at 0x10..0x30 = [12,1,3,maj,min,patch,48,remaining]
    # patch remaining just in case caller changed payload sizes
    u = list(struct.unpack('<8I', out[0x10:0x30]))
    u[6] = CMTD_LEN
    u[7] = total - CMTD_LEN
    out[0x10:0x30] = struct.pack('<8I', *u)
    out += b'mmtd'
    out += struct.pack('<I', total - CMTD_LEN)
    out += struct.pack('<II', meta['mmtd']['flags'], meta['mmtd']['ff'])
    out += struct.pack('<III', *meta['mmtd']['fw'])
    out += struct.pack('<II', meta['mmtd']['entry'], nsect)
    for addr, p in payloads:
        out += b'sect'
        out += struct.pack('<II', addr, len(p))
        out += p
    assert len(out) == total
    if checksum == 'keep':
        final_ck = meta['cmtd']['checksum']
    elif checksum is None:
        final_ck = compute_checksum(bytes(out))
    else:
        final_ck = checksum
    out[8:12] = struct.pack('<I', final_ck)
    return bytes(out)


def headers_to_json(info: dict) -> dict:
    return {
        'cmtd': info['cmtd'],
        'mmtd': {**info['mmtd'],
                 'entry_hex': f"{info['mmtd']['entry']:08X}",
                 'flags_hex': f"{info['mmtd']['flags']:08X}"},
        'sects': [{'index': s['index'], 'off_hex': f"{s['off']:06X}",
                   'addr_hex': f"{s['addr']:08X}", 'len': s['len'],
                   'end_hex': f"{s['end']:08X}"} for s in info['sects']],
    }


def find_sect_by_role(info: dict, role: str) -> dict:
    """role='code' -> largest sect; role='rodata' -> sect holding b'Threshold'
    (fallback: second largest)."""
    sects = info['sects']
    if role == 'code':
        return max(sects, key=lambda s: s['len'])
    if role == 'rodata':
        for s in sects:
            if b'Threshold' in s['payload']:
                return s
        rest = sorted(sects, key=lambda s: s['len'], reverse=True)
        return rest[1] if len(rest) > 1 else rest[0]
    raise ValueError(role)


def code_file_off(info: dict, addr: int) -> int | None:
    for s in info['sects']:
        if s['addr'] <= addr < s['end']:
            return s['off'] + 12 + (addr - s['addr'])
    return None