#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# diff_corpus.py - diff two decompiled corpus indexes across driver versions.
#
# Inputs are `decompiled_index.tsv` files (the merged output of the batch
# decompiler): a header line and one row per function:
#   addr<TAB>file_addr<TAB>size<TAB>name<TAB>file<TAB>line<TAB>status
# `addr` is the Ghidra program address, `file_addr` the file virtual address
# (libcuda+0x...) used everywhere in the analysis.
#
# Matching (function identity across versions):
#   1. Named functions (anything whose name does not start with FUN_) match by
#      name; the same name in both indexes is the same function.
#   2. The address shift between the two builds is estimated as the median of
#      (file_addrB - file_addrA) over name-matched functions.
#   3. Anonymous FUN_* functions match when B has a function at
#      file_addrA + shift; a size difference there is reported as a resize.
#      Functions that do not land on a B function are `removed`; B functions
#      that no A function maps to are `added`.
#
# This is deliberately address/size based (the corpus has no structural hash),
# so it measures exactly what matters for name-map transfer: did the function
# move, grow, disappear or appear.
#
# Usage:
#   diff_corpus.py <indexA.tsv> <indexB.tsv> [--label-a A] [--label-b B]
#                  [--out FILE] [--max-list N] [--json]
#
# Exit status is 0 for a normal report (differences do not fail the tool).
#
# SPDX-License-Identifier: BSD-3-Clause

import argparse
import json
import statistics
import sys
from collections import Counter, defaultdict


def load_index(path):
    funcs = {}
    with open(path) as f:
        f.readline()  # header
        for line in f:
            p = line.rstrip("\n").split("\t")
            if len(p) < 4:
                continue
            try:
                funcs[int(p[1], 16)] = dict(addr=int(p[0], 16), size=int(p[2]),
                                            name=p[3])
            except ValueError:
                continue
    return funcs


def is_named(name):
    return bool(name) and not name.startswith("FUN_")


def label(f, names):
    if f is None:
        return "<none>"
    return f"{names.get(f['name'], f['name'])} (libcuda+0x{f['file']:x}, {f['size']}B)"


def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("index_a")
    ap.add_argument("index_b")
    ap.add_argument("--label-a", default=None)
    ap.add_argument("--label-b", default=None)
    ap.add_argument("--out", default=None, help="write the report here "
                                               "(default: stdout)")
    ap.add_argument("--max-list", type=int, default=40,
                    help="max entries listed per category (default 40)")
    ap.add_argument("--json", action="store_true", help="emit JSON")
    args = ap.parse_args()

    a = load_index(args.index_a)
    b = load_index(args.index_b)
    la = args.label_a or args.index_a
    lb = args.label_b or args.index_b

    named_a = defaultdict(list)
    named_b = defaultdict(list)
    for fa, f in a.items():
        if is_named(f["name"]):
            named_a[f["name"]].append(fa)
    for fb, f in b.items():
        if is_named(f["name"]):
            named_b[f["name"]].append(fb)

    matched_named = []      # (name, fa, fb)
    for name in sorted(set(named_a) & set(named_b)):
        # duplicate symbol names (multiple thunks) match in address order
        for fa, fb in zip(sorted(named_a[name]), sorted(named_b[name])):
            matched_named.append((name, fa, fb))

    deltas = [fb - fa for _, fa, fb in matched_named]
    shift = int(statistics.median(deltas)) if deltas else 0
    dominant = (sum(1 for d in deltas if abs(d - shift) <= 0x40) /
                len(deltas)) if deltas else 0.0

    used_a = set(fa for _, fa, _ in matched_named)
    used_b = set(fb for _, _, fb in matched_named)
    moved = [(n, fa, fb) for n, fa, fb in matched_named if fa != fb]
    resized = [(n, fa, fb) for n, fa, fb in matched_named
               if a[fa]["size"] != b[fb]["size"]]

    shifted = []
    for fa, f in sorted(a.items()):
        if fa in used_a or is_named(f["name"]):
            continue
        fb = fa + shift
        if fb in b and fb not in used_b:
            shifted.append((fa, fb))
            used_a.add(fa)
            used_b.add(fb)
    shifted_resized = [(fa, fb) for fa, fb in shifted
                       if a[fa]["size"] != b[fb]["size"]]
    shifted_ok = [(fa, fb) for fa, fb in shifted
                  if a[fa]["size"] == b[fb]["size"]]

    removed = sorted(fa for fa in a if fa not in used_a)
    added = sorted(fb for fb in b if fb not in used_b)

    anon_a = sum(1 for f in a.values() if not is_named(f["name"]))
    anon_b = sum(1 for f in b.values() if not is_named(f["name"]))
    named_count_a = sum(len(v) for v in named_a.values())
    named_count_b = sum(len(v) for v in named_b.values())
    matched_total = len(matched_named) + len(shifted)

    result = dict(
        label_a=la, label_b=lb,
        total_a=len(a), total_b=len(b),
        named_a=named_count_a, named_b=named_count_b,
        anon_a=anon_a, anon_b=anon_b,
        shift=shift, shift_dominant=round(dominant, 4),
        matched_named=len(matched_named), matched_shifted=len(shifted_ok),
        matched_shifted_resized=len(shifted_resized),
        moved=len(moved), resized_named=len(resized),
        added=len(added), removed=len(removed),
        match_rate_a=round(matched_total / len(a), 4) if a else 0.0,
        match_rate_b=round(matched_total / len(b), 4) if b else 0.0,
        named_transfer_rate=round(len(matched_named) / named_count_a, 4)
        if named_count_a else 0.0,
    )

    def hexlist(addrs):
        return [f"libcuda+0x{x:x} {a[x]['name']} ({a[x]['size']}B)"
                if x in a else f"libcuda+0x{x:x} {b[x]['name']} ({b[x]['size']}B)"
                for x in addrs[:args.max_list]]

    if args.json:
        out = dict(result)
        out["added_list"] = [f"libcuda+0x{x:x} {b[x]['name']} ({b[x]['size']}B)"
                             for x in added[:args.max_list]]
        out["removed_list"] = [f"libcuda+0x{x:x} {a[x]['name']} ({a[x]['size']}B)"
                               for x in removed[:args.max_list]]
        text = json.dumps(out, indent=1)
    else:
        lines = []
        lines.append(f"corpus diff: {la}  ->  {lb}")
        lines.append("=" * 72)
        lines.append(f"functions        : A={len(a)}  B={len(b)}  "
                     f"(named A={len(named_a)} B={len(named_b)}; "
                     f"anonymous A={anon_a} B={anon_b})")
        lines.append(f"address shift    : 0x{shift:x} "
                     f"(dominant for {dominant*100:.1f}% of name matches)")
        lines.append(f"matched          : {matched_total} "
                     f"({len(matched_named)} by name, "
                     f"{len(shifted_ok)} anonymous at shifted address)")
        lines.append(f"match rate       : A {result['match_rate_a']*100:.1f}%  "
                     f"B {result['match_rate_b']*100:.1f}%  "
                     f"(name map transfer {result['named_transfer_rate']*100:.1f}%)")
        lines.append(f"moved            : {len(moved)} named functions changed "
                     f"address")
        lines.append(f"resized          : {len(resized)} named + "
                     f"{len(shifted_resized)} anonymous matched functions "
                     f"changed size")
        lines.append(f"added            : {len(added)}")
        lines.append(f"removed          : {len(removed)}")
        lines.append("")
        if moved:
            lines.append(f"-- moved named (first {args.max_list}) --")
            for n, fa, fb in moved[:args.max_list]:
                lines.append(f"   {n}: libcuda+0x{fa:x} -> libcuda+0x{fb:x} "
                             f"({fb-fa:+#x})")
        if resized:
            lines.append(f"-- resized named (first {args.max_list}) --")
            for n, fa, fb in resized[:args.max_list]:
                lines.append(f"   {n}: libcuda+0x{fa:x} {a[fa]['size']}B -> "
                             f"{b[fb]['size']}B")
        if shifted_resized:
            lines.append(f"-- resized anonymous (first {args.max_list}) --")
            for fa, fb in shifted_resized[:args.max_list]:
                lines.append(f"   libcuda+0x{fa:x} {a[fa]['size']}B -> "
                             f"libcuda+0x{fb:x} {b[fb]['size']}B")
        if added:
            lines.append(f"-- added (first {args.max_list}) --")
            lines.extend("   " + x for x in
                         hexlist(added))
        if removed:
            lines.append(f"-- removed (first {args.max_list}) --")
            lines.extend("   " + x for x in
                         hexlist(removed))
        text = "\n".join(lines)

    if args.out:
        with open(args.out, "w") as f:
            f.write(text + "\n")
    else:
        print(text)
    return 0


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