#!/usr/bin/env python3

# This script takes data generated by the profiler in Astral and creates a flamegraph out of it.
# Format of the profiler file:
# 1 byte indicating the number of calls in the backtrace
# 8-byte Pointers starting by deepest call in the stack
# Repeat until EOF

import sys
import os
import subprocess

if len(sys.argv) != 3:
    print(sys.argv[0] + ": usage: " + sys.argv[0] + " INPUT_FILE OUTPUT_FILE")
    exit(1)

# Open needed files
in_file = open(sys.argv[1], "rb")
out_file = open(sys.argv[2], "w")

# Download the flamegraph script if nescessary
os.chdir(os.path.dirname(__file__))

if not os.access("flamegraph.pl", os.R_OK):
    os.system("wget https://raw.githubusercontent.com/brendangregg/FlameGraph/refs/heads/master/flamegraph.pl")
    os.system("chmod +x flamegraph.pl")

print("Starting subprocesses")

# Start the flamegraph renderer
proc_fg = subprocess.Popen(
    ["./flamegraph.pl"],
    stdin=subprocess.PIPE,
    stdout=out_file,
    text=True
)

traces = {}

mappings = {}

print("Getting function information")
# get info out of addr2line
while True:
    backtrace_size = in_file.read(1)
    if not backtrace_size:
        # EOF
        break

    backtrace_size = int.from_bytes(backtrace_size, byteorder='little', signed=False)
    backtrace = []

    # Get backtrace, calling addr2line for non-cached mappings
    i = backtrace_size
    while i > 0:
        i -= 1

        addr = in_file.read(8)
        if not addr:
            print(("Corrupted file (unexpected EOF, reading %d entries)" % (backtrace_size)), file=sys.stderr)
            exit(1)

        addr = hex(int.from_bytes(addr, byteorder='little', signed=False))

        # Add it to the backtrace
        backtrace.append(addr)
        if addr in mappings:
            continue

        # Get the mapping
        msg = addr + "\n"
        # Start and send the data to addr2line
        data = subprocess.Popen(
                ["addr2line", "-i", "-s", "-f", "-e", "../build-x86_64/iso/astral"],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                text=True
        ).communicate(input=msg)[0].split('\n')

        mappings[addr] = data

    # format the response in a way flamegraph expects
    formatted_str = ""
    for addr in backtrace:
        map_data = mappings[addr]

        aux_function = ""
        for data_line in map_data:
            if aux_function == "":
                aux_function = data_line
                continue

            file, line = data_line.split(":")

            formatted_str = file + "`" + aux_function + ";" + formatted_str
            aux_function = ""

    formatted_str = formatted_str[:-1] # remove ;

    if not formatted_str in traces:
        traces[formatted_str] = 0

    traces[formatted_str] += 1

print("Generating flamegraph")
# send the trace over to flamegraph
for trace in traces:
    final_string = trace + " " + str(traces[trace])
    proc_fg.stdin.write(final_string + "\n")
    proc_fg.stdin.flush()

print("Waiting...");
# wait until it is done with its thing
proc_fg.stdin.close()
proc_fg.wait()