import gdb
import os


def libpoplar_base():
    pid = gdb.selected_inferior().pid
    with open("/proc/%d/maps" % pid) as mappings:
        for line in mappings:
            if "libpoplar.so" in line and "r-xp" in line:
                start = int(line.split("-", 1)[0], 16)
                offset = int(line.split()[2], 16)
                return start - offset
    raise gdb.GdbError("libpoplar.so is not mapped")


class Capture(gdb.Breakpoint):
    def __init__(self, address, output):
        super().__init__("*0x%x" % address, internal=False)
        self.output = output
        self.call = 0
        self.selected_call = int(os.environ.get("IPU_PREFIX_CALL", "1"), 0)

    def stop(self):
        self.call += 1
        if self.call != self.selected_call:
            return False
        self.enabled = False
        loader = int(gdb.parse_and_eval("$rbx"))
        inferior = gdb.selected_inferior()
        pointer_size = gdb.lookup_type("void").pointer().sizeof
        buffer_address = int.from_bytes(
            inferior.read_memory(loader + 0x10, pointer_size),
            byteorder="little")
        capture_bytes = int(os.environ.get("IPU_PREFIX_BYTES", "8192"), 0)
        prefix = inferior.read_memory(buffer_address, capture_bytes).tobytes()
        with open(self.output, "wb") as output:
            output.write(prefix)
        gdb.write("captured buffer prefix at %#x in %s\n" %
                  (buffer_address, self.output))
        return False


base = libpoplar_base()
output = os.environ.get("IPU_PREFIX_CAPTURE", "/tmp/ipu-batch-prefix.bin")
Capture(base + 0x3BD1749, output)