#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# annotate_corpus.py - turn the raw decompiled corpus into an annotated,
# human-navigable corpus.
#
# For every function in the decompiled buckets it rewrites the header comment
# with:
#   * the recovered name (from a names.tsv map) or the Ghidra FUN_ name;
#   * file virtual address and size;
#   * direct callers and callees (from callgraph.tsv), labelled with names;
#   * a "seen in" marker for functions referenced by the docs (optional).
# It also substitutes recovered names for FUN_<addr> references inside the
# function bodies, so a renamed function is readable everywhere.
#
# Inputs (produced by the toolchain):
#   <outdir>/buckets/bucket_*.c          raw decompiled C
#   <outdir>/decompiled_index.tsv        addr file_addr size name file line status
#   <indexdir>/callgraph.tsv             caller callee caller_file callee_file
#   <names.tsv>                          file_vaddr<TAB>name
#
# Usage:
#   annotate_corpus.py <outdir> <indexdir> <names.tsv> <annotated_outdir>
#
# The output directory is gitignored like the raw corpus; this tool is tracked
# and reusable. Addresses are file virtual addresses in all inputs; the corpus
# uses Ghidra program addresses, so the tool translates via the index's
# addr/file_addr columns (no hard-coded image base).

import os
import re
import sys
from collections import defaultdict


def load_index(path):
    rows = {}
    with open(path) as f:
        header = f.readline()
        for line in f:
            p = line.rstrip("\n").split("\t")
            if len(p) < 6:
                continue
            prog = int(p[0], 16)
            rows[prog] = dict(file=int(p[1], 16), size=int(p[2]),
                              name=p[3], file_name=p[4], line=int(p[5]))
    return rows


def load_callgraph(path):
    callees = defaultdict(list)
    callers = defaultdict(list)
    with open(path) as f:
        f.readline()
        for line in f:
            p = line.rstrip("\n").split("\t")
            if len(p) < 4:
                continue
            cf, kf = int(p[2], 16), int(p[3], 16)
            callees[cf].append(kf)
            callers[kf].append(cf)
    return callees, callers


def load_names(path):
    names = {}
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            p = line.split("\t")
            if len(p) >= 2:
                try:
                    names[int(p[0], 16)] = p[1]
                except ValueError:
                    pass
    return names


def label(file_addr, names, index):
    name = names.get(file_addr)
    if name:
        return f"{name} (libcuda+0x{file_addr:x})"
    return f"libcuda+0x{file_addr:x}"


def main():
    if len(sys.argv) < 5:
        print(__doc__)
        sys.exit(2)
    outdir, indexdir, namemap, out = sys.argv[1:5]
    index = load_index(os.path.join(outdir, "decompiled_index.tsv"))
    callees, callers = load_callgraph(os.path.join(indexdir, "callgraph.tsv"))
    names = load_names(namemap)
    # program address -> recovered name, for body substitution
    prog_names = {p: names[r["file"]] for p, r in index.items() if r["file"] in names}

    os.makedirs(os.path.join(out, "buckets"), exist_ok=True)
    n = 0
    for bucket in sorted(os.listdir(os.path.join(outdir, "buckets"))):
        if not bucket.endswith(".c"):
            continue
        src = os.path.join(outdir, "buckets", bucket)
        dst = os.path.join(out, "buckets", bucket)
        with open(src) as f:
            text = f.read()

        def annotate(m):
            nonlocal n
            prog = int(m.group("addr"), 16)
            file_addr = int(m.group("file"), 16)
            size = int(m.group("size"))
            rec = index.get(prog, {})
            name = names.get(file_addr, rec.get("name", f"FUN_{prog:08x}"))
            cl = sorted(set(callees.get(file_addr, [])))
            cr = sorted(set(callers.get(file_addr, [])))
            lines = [
                "/* ------------------------------------------------------------------",
                f" * {name}",
                f" * file_addr=libcuda+0x{file_addr:x} size={size}",
            ]
            if cr:
                shown = [label(a, names, index) for a in cr[:8]]
                more = f" (+{len(cr)-8} more)" if len(cr) > 8 else ""
                lines.append(" * called by: " + "; ".join(shown) + more)
            if cl:
                shown = [label(a, names, index) for a in cl[:8]]
                more = f" (+{len(cl)-8} more)" if len(cl) > 8 else ""
                lines.append(" * calls:     " + "; ".join(shown) + more)
            lines.append(" *----------------------------------------------------------------- */")
            n += 1
            return "\n" + "\n".join(lines) + "\n"

        text = re.sub(
            r"\n/\* -+\n \* (?P<name>[^\n]*)\n \* addr=(?P<addr>0x[0-9a-f]+) "
            r"file_addr=libcuda\+(?P<file>0x[0-9a-f]+) size=(?P<size>\d+)\n"
            r" \*-+ \*/",
            annotate, text)

        if prog_names:
            def sub(m):
                a = int(m.group(1), 16)
                return prog_names.get(a, m.group(0))
            text = re.sub(r"\bFUN_([0-9a-f]{8})\b", sub, text)

        with open(dst, "w") as f:
            f.write(text)

    print(f"[annotate_corpus] annotated {n} functions -> {out}/buckets")
    print(f"[annotate_corpus] names applied: {len(names)}")


if __name__ == "__main__":
    main()