#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
#
# subsystem_digest.py - per-subsystem call-tree digest from the tracked
# subsystem annotations plus the callgraph index.
#
# For every subsystem in subsystems.tsv it prints:
#   * members (address, name, role) - the annotated functions;
#   * entry functions - members that are not called by another member of the
#     same subsystem (the roots of the subsystem's call tree);
#   * a depth-limited call tree from each entry, with the direct callees of a
#     node grouped by their subsystem;
#   * a cross-subsystem edge summary (how many calls each subsystem makes into
#     every other one) - the fastest way to see the layering.
#
# The digest is a navigation aid: it shows where to start reading the
# decompiled corpus for a subsystem and which other subsystems it touches,
# without opening Ghidra. Addresses are FILE virtual addresses (libcuda+0x...).
#
# Usage:
#   subsystem_digest.py <subsystems.tsv> <callgraph.tsv> [--out FILE]
#                       [--depth N] [--max-per-node N] [--subsystem NAME]
#
# <subsystems.tsv> format (comments allowed):
#   file_vaddr<TAB>name<TAB>subsystem<TAB>role<TAB>evidence
# <callgraph.tsv> format (header):
#   caller<TAB>callee<TAB>caller_file<TAB>callee_file
#
# SPDX-License-Identifier: BSD-3-Clause

import argparse
import sys
from collections import defaultdict


def load_subsystems(path):
    members = {}
    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) < 4:
                continue
            try:
                addr = int(p[0], 16)
            except ValueError:
                continue
            members[addr] = dict(name=p[1], subsystem=p[2],
                                 role=p[3],
                                 evidence=p[4] if len(p) > 4 else "")
    return members


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


def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("subsystems")
    ap.add_argument("callgraph")
    ap.add_argument("--out", default=None,
                    help="write the digest here (default: stdout)")
    ap.add_argument("--depth", type=int, default=2,
                    help="call-tree depth below each entry (default 2)")
    ap.add_argument("--max-per-node", type=int, default=8,
                    help="max callees shown per node (default 8)")
    ap.add_argument("--subsystem", default=None,
                    help="only digest this subsystem")
    args = ap.parse_args()

    members = load_subsystems(args.subsystems)
    callees, callers = load_callgraph(args.callgraph)

    by_sub = defaultdict(dict)
    for addr, m in members.items():
        by_sub[m["subsystem"]][addr] = m

    def name_of(a):
        m = members.get(a)
        if m:
            return f"{m['name']} (libcuda+0x{a:x})"
        return f"libcuda+0x{a:x}"

    def sub_of(a):
        m = members.get(a)
        return m["subsystem"] if m else "unknown"

    out = []
    out.append("# Subsystem call-tree digest (generated by tools/subsystem_digest.py)")
    out.append(f"# subsystems.tsv: {args.subsystems}")
    out.append(f"# callgraph.tsv : {args.callgraph}")
    out.append("# Addresses are FILE virtual addresses (libcuda+0x...).")
    out.append("# entry = annotated member with no caller inside its own subsystem.")
    out.append("")

    # cross-subsystem edge counts
    cross = defaultdict(int)
    for addr, m in members.items():
        for c in callees.get(addr, ()):
            cross[(m["subsystem"], sub_of(c))] += 1

    selected = sorted(by_sub)
    if args.subsystem:
        selected = [s for s in selected if s == args.subsystem]
        if not selected:
            print(f"error: subsystem {args.subsystem!r} not in {args.subsystems}",
                  file=sys.stderr)
            return 1

    for sub in selected:
        ms = by_sub[sub]
        out.append("=" * 72)
        out.append(f"## {sub} ({len(ms)} annotated members)")
        out.append("=" * 72)
        out.append("members:")
        for a in sorted(ms):
            m = ms[a]
            out.append(f"  libcuda+0x{a:x}\t{m['name']}\t{m['role']}")
        entries = sorted(a for a in ms
                         if not (callers.get(a, set()) & set(ms)))
        if not entries:
            entries = sorted(ms)
            out.append("entries: (no entry: every member is called from inside "
                       "the subsystem; showing all)")
        else:
            out.append("entries:")
        for a in entries:
            out.append(f"  {name_of(a)}: {ms[a]['role']}")
        out.append("call tree:")

        def walk(a, depth, indent, seen):
            kids = sorted(callees.get(a, ()))
            if not kids or depth <= 0:
                return
            groups = defaultdict(list)
            for k in kids:
                groups[sub_of(k)].append(k)
            for g in sorted(groups):
                shown = groups[g][:args.max_per_node]
                more = f" (+{len(groups[g]) - args.max_per_node} more)" \
                    if len(groups[g]) > args.max_per_node else ""
                out.append(f"{indent}-> {g}: " +
                           "; ".join(name_of(k) for k in shown) + more)
            for k in kids:
                if k in seen:
                    continue
                seen.add(k)
                walk(k, depth - 1, indent + "   ", seen)

        for a in entries:
            out.append(f"  {name_of(a)}")
            walk(a, args.depth, "    ", {a})
        out.append("")

    out.append("=" * 72)
    out.append("## cross-subsystem calls (edges from -> to)")
    out.append("=" * 72)
    for (src, dst) in sorted(cross):
        out.append(f"  {src:<16} -> {dst:<16} {cross[(src, dst)]}")

    text = "\n".join(out)
    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())