#!/usr/bin/env python3
import argparse
import subprocess
import sys


def get_syms(program):
    result = subprocess.run(
        ["nm", "-defined-only", "-demangle", program],
        check=True,
        capture_output=True,
        encoding="utf8",
    )
    syms = {}
    for line in result.stdout.split("\n"):
        if not line.strip():
            continue
        addr, _, name = line.split(" ", 2)
        if name.startswith("ivy_builtin$"):
            name = name.split("$", 2)[1]
        syms[addr] = name
    return syms


def handle_trace(lines, syms={}, only_debug=False):
    objs = {}

    def obj(x):
        if int(x, 16) % 2 == 1:
            return int(x, 16) >> 1
        if x in objs:
            return objs[x]
        elif x in syms:
            return syms[x]
        return x

    for line in lines:
        if not line.strip():
            print("Blank line")
            continue
        kind, *args = line.split()
        if kind == "MAKE":
            objs[args[0]] = syms.get(args[1], f"fn@{args[1]}")
            if not only_debug:
                print("MAKE", objs[args[0]])
        elif kind == "COPY":
            objs[args[1]] = obj(args[0])
        elif kind == "APP":
            if int(args[1], 16) == 0:
                objs[args[0]] = f"({obj(args[0])})"
            else:
                objs[args[0]] = f"({obj(args[0])} {obj(args[1])})"
            if not only_debug:
                print("APP", objs[args[0]])
        elif kind == "DEBUG":
            if only_debug:
                print(obj(args[0]))
            else:
                print("DEBUG", obj(args[0]))
        elif kind == "UPD8":
            pass  # who cares :P
        else:
            if not only_debug:
                print(kind, *(obj(x) for x in args))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Make an Ivy runtime trace more readable",
    )
    parser.add_argument("program", help="The program that produced the trace")
    parser.add_argument("trace", help="A file containing the trace")
    parser.add_argument(
        "-d",
        "--only-debug",
        action="store_true",
        help="If set, only print the DEBUG outputs",
    )
    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(0)
    args = parser.parse_args()

    syms = get_syms(args.program)
    with open(args.trace, "r") as f:
        lines = [line.strip() for line in f if line.strip()]
    handle_trace(lines, syms, ags.only_debug)