#!/usr/bin/env python3
"""fsr2-tierdef-audit — exhaustive per-channel comparison of two MOZA captures'
tier-definitions (e.g. plugin vs PitHouse), mapped to real channel names.

For each capture it:
  * reassembles h2b + b2h session-0xNN data chunks (7E-framed wire frames from
    tools/pcap_to_jsonl.py), de-duplicating retransmits by seq per session;
  * decodes the wheel's channel catalog (b2h sess 0x02) into a 1-based
    idx -> URL map, implementing the FULL record grammar used by
    Telemetry/Frames/ChannelCatalogParser.cs (full-URL, 0x01 / \\1 / \\p
    abbreviations with \\t \\P \\b corner expansions, and back-references);
  * decodes the host's tier-definition (h2b sess 0x02) into per-channel
    (comp_code, bit_width, tier_flag) using the V2Type02 TLV layout
    (Telemetry/Frames/TierDefinitionBuilder.cs);

then diffs the two captures channel-by-channel and classifies every difference.

Usage:
    tools/fsr2-tierdef-audit <plugin.jsonl> <pithouse.jsonl>
    tools/fsr2-tierdef-audit <plugin.jsonl> <pithouse.jsonl> --telemetry Data/Telemetry.json
"""
from __future__ import annotations
import argparse
import json
import re
import struct
import sys
from collections import defaultdict
from pathlib import Path

# ── compression code table (mirror of Telemetry/Protocol/CompressionTable.cs) ──
# (code, bit_width) -> name.  Code 0x16 is overloaded (tyre_pressure_1 @12,
# brake_temp_1 @16); 0x0D and 0x14 each map several signed/unsigned names — we
# key by (code,bw) where it disambiguates, else fall back to code alone.
COMP_BY_CODE_BW = {
    (0x00, 1): "bool",
    (0x02, 8): "int8_t",
    (0x01, 8): "uint8_t",
    (0x04, 16): "uint16_t",
    (0x05, 16): "int16_t",
    (0x07, 32): "float",
    (0x0D, 5): "int30/uint31",
    (0x0E, 10): "percent_1",
    (0x17, 10): "float_001",
    (0x0F, 16): "float_6000_1",
    (0x15, 16): "float_600_2",
    (0x16, 12): "tyre_pressure_1",
    (0x11, 14): "tyre_temp_1",
    (0x12, 14): "track_temp_1",
    (0x13, 14): "oil_pressure_1",
    (0x16, 16): "brake_temp_1",
    (0x14, 4): "uint3/uint8",
    (0x03, 4): "uint15",
    (0x08, 32): "int32_t",
    (0x09, 32): "uint32_t",
    (0x18, 24): "uint24_t",
    (0x0A, 64): "double",
    (0x0B, 64): "location_t",
    (0x0C, 64): "int64_t",
    (0x19, 64): "uint64_t",
}
COMP_BY_CODE = {
    0x00: "bool", 0x01: "uint8_t", 0x02: "int8_t", 0x03: "uint15", 0x04: "uint16_t",
    0x05: "int16_t", 0x07: "float", 0x08: "int32_t", 0x09: "uint32_t", 0x0A: "double",
    0x0B: "location_t", 0x0C: "int64_t", 0x0D: "int30/uint31", 0x0E: "percent_1",
    0x0F: "float_6000_1", 0x11: "tyre_temp_1", 0x12: "track_temp_1", 0x13: "oil_pressure_1",
    0x14: "uint3/uint8", 0x15: "float_600_2", 0x16: "tyre_pressure_1/brake_temp_1",
    0x17: "float_001", 0x18: "uint24_t", 0x19: "uint64_t",
}

def comp_name(code, bw):
    return COMP_BY_CODE_BW.get((code, bw)) or COMP_BY_CODE.get(code) or f"0x{code:X}?"


# ── wire reassembly ────────────────────────────────────────────────────────
def load(path):
    return [json.loads(l) for l in open(path) if l.strip()]

def reassemble(path):
    """Return (h2b_streams, b2h_streams): session -> reassembled data bytes.

    Frames are 7E-framed: [7E][len][group][dev][payload..][csum].
    Session data sub-frame inside group 0x43 (h2b) / 0xC3 dev 0x71 (b2h):
        [7C][00][session][stype][seq_lo][seq_hi][data..]
    only stype 0x01 (data) chunks are reassembled, deduped by seq (keep first).
    """
    h2b = defaultdict(dict)
    b2h = defaultdict(dict)
    for r in load(path):
        b = bytes.fromhex(r["hex"])
        if len(b) < 6 or b[0] != 0x7E:
            continue
        group, dev = b[2], b[3]
        pl = b[4:4 + b[1]]
        if len(pl) < 6 or pl[0] != 0x7C or pl[1] != 0x00 or pl[3] != 0x01:
            continue
        session = pl[2]
        seq = pl[4] | (pl[5] << 8)
        # Each data chunk ends with a 4-byte CRC (frame csum already excluded
        # from pl). Strip it so concatenated chunks form a clean TLV stream —
        # mirrors moza_trace's raw[10:-5]. Tiny chunks (<=4 data bytes) carry
        # no CRC.
        data = pl[6:-4] if len(pl) > 10 else pl[6:]
        if r["dir"] == "h2b" and group == 0x43:
            tbl = h2b
        elif r["dir"] == "b2h" and group == 0xC3 and dev == 0x71:
            tbl = b2h
        else:
            continue
        if seq not in tbl[session]:
            tbl[session][seq] = data
    merge = lambda d: {s: b"".join(c[k] for k in sorted(c)) for s, c in d.items()}
    return merge(h2b), merge(b2h)


# ── catalog decode (mirror of ChannelCatalogParser.TryParse) ───────────────
def decode_catalog(buf):
    """idx(1-based) -> URL. Implements full-URL, 0x01/\\1/\\p abbrevs, back-refs."""
    parsed = {}
    i = 0
    n = len(buf)
    while i + 6 < n:
        tag = buf[i]
        if tag == 0x06:  # END marker 06 04 00 00 00 <u32>
            if i + 8 < n and buf[i+1] == 0x04 and buf[i+2] == 0 and buf[i+3] == 0 and buf[i+4] == 0:
                i += 9
                continue
        if tag != 0x04:
            i += 1
            continue
        param = struct.unpack_from("<I", buf, i + 1)[0]
        if param < 1 or param >= 200 or i + 5 + param > n:
            i += 1
            continue
        idx = buf[i + 5]
        url_len = param - 1
        url_start = i + 6
        if url_len == 0:  # back-ref: resolve from already-parsed idx
            if idx in parsed:
                pass  # binding already known; keep it
            i += 5 + param
            continue
        first = buf[url_start]
        body = buf[url_start:url_start + url_len]
        url = None
        if first == 0x01:
            url = "v1/gameData/" + body[1:].decode("latin1")
        elif first == 0x5C and url_len >= 2 and body[1] == 0x31:  # \1
            suffix = body[2:].decode("latin1")
            suffix = (suffix.replace("\\t", "TyreTemp").replace("\\P", "TyrePressure")
                      .replace("\\b", "BrakeTemp").replace("{FL}", "FrontLeft")
                      .replace("{FR}", "FrontRight").replace("{RL}", "RearLeft")
                      .replace("{RR}", "RearRight"))
            url = "v1/gameData/" + suffix
        elif first == 0x5C and url_len >= 2 and body[1] == 0x70:  # \p
            url = "v1/gameData/patch/" + body[2:].decode("latin1")
        elif body[:3] == b"v1/" or body[:3] == b"v0/":
            url = body.decode("latin1")
        else:
            i += 1
            continue
        # validate printable
        if url is None or any(not (0x20 <= ord(c) <= 0x7E) for c in url):
            i += 1
            continue
        if idx >= 1:
            parsed[idx] = url
        i += 5 + param
    return parsed


# ── tier-def decode (V2Type02 TLV) ─────────────────────────────────────────
def decode_tierdef(buf):
    """Return dict: idx -> {'comp':code,'bw':bits,'tiers':set(relative_tier)}.

    The host re-broadcasts the full tier set several times with an incrementing
    flag base; we group TLV records into emissions at each END marker, normalise
    each tier's flag to a relative sub-tier (flag - emission_base), and collect
    per channel its (comp,bw) plus the set of relative sub-tiers it appears in.
    A channel's (comp,bw) MUST be consistent wherever it appears; we record all
    seen values to surface any inconsistency.
    """
    chans = defaultdict(lambda: {"comp": set(), "bw": set(), "tiers": set()})
    i = 0
    n = len(buf)
    cur_emission_tiers = []  # list of (flag, [(idx,comp,bw)])
    def flush(emission):
        if not emission:
            return
        base = min(f for f, _ in emission)
        for flag, recs in emission:
            rel = flag - base
            for idx, comp, bw in recs:
                chans[idx]["comp"].add(comp)
                chans[idx]["bw"].add(bw)
                chans[idx]["tiers"].add(rel)
    while i + 5 <= n:
        tag = buf[i]
        size = struct.unpack_from("<I", buf, i + 1)[0]
        if i + 5 + size > n:
            break
        data = buf[i + 5:i + 5 + size]
        if tag == 0x01 and size >= 1:  # TIER: flag + 16B records
            flag = data[0]
            recs = []
            body = data[1:]
            for j in range(0, len(body) - 15, 16):
                idx, comp, bw, _res = struct.unpack_from("<IIII", body, j)
                if 0 < idx < 2000 and 0 < bw <= 64:
                    recs.append((idx, comp, bw))
            cur_emission_tiers.append((flag, recs))
        elif tag == 0x06:  # END marker -> emission boundary
            flush(cur_emission_tiers)
            cur_emission_tiers = []
        i += 5 + size
    flush(cur_emission_tiers)
    return chans


def telemetry_compressions(path):
    """url -> compression string from Telemetry.json."""
    if not path or not Path(path).exists():
        return {}
    raw = open(path, encoding="utf-8").read()
    out = {}
    # each channel object has a "url" and a sibling "compression"
    for m in re.finditer(r'"(v[01]/[^"]+)"', raw):
        url = m.group(1)
        seg = raw[m.start():m.start() + 600]
        cm = re.search(r'"compression"\s*:\s*"([^"]+)"', seg)
        if cm and url not in out:
            out[url] = cm.group(1)
    return out


def summarise(name, h2b, b2h):
    cat = decode_catalog(b2h.get(0x02, b""))
    td = decode_tierdef(h2b.get(0x02, b""))
    print(f"[{name}] catalog idx->url = {len(cat)} entries; tier-def channels = {len(td)}")
    return cat, td


def fmt(v):
    return sorted(v) if len(v) != 1 else next(iter(v))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("plugin")
    ap.add_argument("pithouse")
    ap.add_argument("--telemetry", default="Data/Telemetry.json")
    args = ap.parse_args()

    ph_h2b, ph_b2h = reassemble(args.plugin)
    rh_h2b, rh_b2h = reassemble(args.pithouse)
    cat_p, td_p = summarise("PLUGIN", ph_h2b, ph_b2h)
    cat_h, td_h = summarise("PITHOUSE", rh_h2b, rh_b2h)
    tj = telemetry_compressions(args.telemetry)

    # merged catalog (prefer PitHouse; fall back to plugin)
    cat = dict(cat_p)
    cat.update(cat_h)

    all_idx = sorted(set(td_p) | set(td_h))
    print(f"\nchannels in either tier-def: {len(all_idx)}  "
          f"(plugin={len(td_p)}, pithouse={len(td_h)}, catalog names={len(cat)})")

    classes = defaultdict(list)
    rows = []
    for idx in all_idx:
        url = cat.get(idx, f"<idx {idx}: no catalog name>")
        p = td_p.get(idx)
        h = td_h.get(idx)
        pj = tj.get(url, "")
        if p and not h:
            cls = "PLUGIN_ONLY"
        elif h and not p:
            cls = "OMITTED_BY_PLUGIN"
        else:
            pc, pb = fmt(p["comp"]), fmt(p["bw"])
            hc, hb = fmt(h["comp"]), fmt(h["bw"])
            if pc != hc:
                cls = "COMP_MISMATCH"
            elif pb != hb:
                cls = "BW_MISMATCH"
            elif p["tiers"] != h["tiers"]:
                cls = "TIER_ONLY_DIFF"
            else:
                cls = "IDENTICAL"
        classes[cls].append(idx)
        rows.append((idx, url, p, h, pj, cls))

    print("\n=== classification summary ===")
    for cls in ["COMP_MISMATCH", "BW_MISMATCH", "OMITTED_BY_PLUGIN",
                "PLUGIN_ONLY", "TIER_ONLY_DIFF", "IDENTICAL"]:
        print(f"  {cls:18s}: {len(classes[cls])}")

    def comp_str(rec):
        if not rec:
            return "-"
        c, b = fmt(rec["comp"]), fmt(rec["bw"])
        cc = c if isinstance(c, int) else c
        nm = comp_name(c, b) if isinstance(c, int) and isinstance(b, int) else f"{c}/{b}"
        return f"{nm}(0x{c:X}/{b})" if isinstance(c, int) else f"{c}/{b}"

    for cls in ["COMP_MISMATCH", "BW_MISMATCH", "OMITTED_BY_PLUGIN", "PLUGIN_ONLY", "TIER_ONLY_DIFF"]:
        idxs = classes[cls]
        if not idxs:
            continue
        print(f"\n=== {cls} ({len(idxs)}) ===")
        for idx in idxs:
            _, url, p, h, pj, _ = next(r for r in rows if r[0] == idx)
            short = url.replace("v1/gameData/", "")
            print(f"  idx {idx:3d}  {short:34s}  plugin={comp_str(p):28s}  pithouse={comp_str(h):28s}  tj={pj}")


if __name__ == "__main__":
    main()
