#!/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 addr2line TODO: other architectures
proc_a2l = subprocess.Popen(
        ["addr2line", "-s", "-f", "-e", "../build-x86_64/iso/astral"],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        text=True
)

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

traces = {}

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)

    # Send the addresses to addr2line
    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))

        proc_a2l.stdin.write(addr + "\n")
        proc_a2l.stdin.flush()

    # Read back the lines from addr2line and format them in a way flamegraph expects
    i = backtrace_size
    formatted_str = ""
    while i > 0:
        i -= 1

        function = proc_a2l.stdout.readline().replace('\n', '')
        file, line = proc_a2l.stdout.readline().replace('\n', '').split(":")

        formatted_str = file + "`" + function + ";" + formatted_str

    formatted_str = formatted_str[:-1] # remove ;

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

    traces[formatted_str] += 1
        
proc_a2l.stdin.close()

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()