"""Decode Tran across multiple segments by resetting at 40 02 headers."""
import sys
sys.path.insert(0, ".")
from analysis.load_bundle import _parse_txt
from minimateplus.waveform_codec import walk_body, find_data_start


def s4(n):
    return n if n < 8 else n - 16


def i8(b):
    return b if b < 128 else b - 256


def decode_full_tran(body):
    """Decode all Tran samples in the body, walking through segments."""
    if len(body) < 7 or body[0:3] != b"\x00\x02\x00":
        return None
    T0 = int.from_bytes(body[3:5], "big", signed=True)
    T1 = int.from_bytes(body[5:7], "big", signed=True)

    # Locate first tag
    i = 7
    while i + 1 < len(body) and body[i] not in (0x00, 0x10, 0x20, 0x30, 0x40):
        i += 1

    blocks = walk_body(body, i)
    T = [T0, T1]
    cur = T1
    for bi, blk in enumerate(blocks):
        if blk.tag_hi == 0x40:
            # Segment header — try interpreting bytes [0:2] as new T anchor
            if len(blk.data) >= 2:
                new_anchor = int.from_bytes(blk.data[0:2], "big", signed=True)
                # The next sample IS this anchor value, NOT a delta from cur.
                T.append(new_anchor)
                cur = new_anchor
        elif blk.tag_hi == 0x10:
            for byte in blk.data:
                for nib in ((byte >> 4) & 0xF, byte & 0xF):
                    cur += s4(nib)
                    T.append(cur)
        elif blk.tag_hi == 0x20:
            for byte in blk.data:
                cur += i8(byte)
                T.append(cur)
        elif blk.tag_hi == 0x00:
            # RLE: append NN zero deltas
            for _ in range(blk.tag_lo):
                T.append(cur)
        # 30 NN: skip
    return T


def main():
    for stem in ("M529LL1L.V70", "M529LL1L.JQ0", "M529LL1A.SP0", "M529LL1A.SS0", "M529LL1A.SV0"):
        path = f"tests/fixtures/5-11-26/{stem}"
        with open(path, "rb") as f:
            body = f.read()[43:-26]
        _, samples = _parse_txt(path + ".TXT")
        truth_T = [round(v*200) for v in samples["Tran"]]
        n_truth = len(truth_T)

        decoded = decode_full_tran(body)
        n = min(len(decoded), n_truth)
        matches = sum(1 for i in range(n) if decoded[i] == truth_T[i])
        # Find first divergence
        div_at = -1
        for i in range(n):
            if decoded[i] != truth_T[i]:
                div_at = i
                break
        print(f"{stem}: decoded={len(decoded)}, truth={n_truth}, matches={matches}/{n}, first div={div_at}")
        if div_at >= 0 and div_at < 30:
            print(f"  truth around div [{max(0,div_at-3)}:{div_at+8}]: {truth_T[max(0,div_at-3):div_at+8]}")
            print(f"  pred  around div [{max(0,div_at-3)}:{div_at+8}]: {decoded[max(0,div_at-3):div_at+8]}")


if __name__ == "__main__":
    main()