/*
 * Copyright (C) 2025, Pierrick Bouvier <pierrick.bouvier@linaro.org>
 *
 * Generates a trace compatible with uftrace (similar to uftrace record).
 * https://github.com/namhyung/uftrace
 *
 * See docs/about/emulation.rst|Uftrace for details and examples.
 *
 * SPDX-License-Identifier: GPL-2.0-or-later
 */

#include <qemu-plugin.h>
#include <glib.h>
#include <stdio.h>
#include <sys/stat.h>
#include <sys/time.h>
#include <time.h>
#include <unistd.h>

#define MiB     (INT64_C(1) << 20)
#define NANOSECONDS_PER_SECOND 1000000000LL
#define TRACE_FLUSH_SIZE (32 * MiB)
#define TRACE_ID_SCALE 100

QEMU_PLUGIN_EXPORT int qemu_plugin_version = QEMU_PLUGIN_VERSION;

typedef struct {
    GArray *s;
} Callstack;

typedef struct {
    uint64_t pc;
    uint64_t frame_pointer;
} CallstackEntry;

typedef struct {
    GArray *t;
    GString *path;
    GString *name;
    uint32_t id;
} Trace;

typedef struct Cpu Cpu;

typedef struct {
    void (*init)(Cpu *cpu);
    void (*end)(Cpu *cpu);
    uint64_t (*get_frame_pointer)(Cpu *cpu);
    uint8_t (*get_privilege_level)(Cpu *cpu);
    uint8_t (*num_privilege_levels)(void);
    const char *(*get_privilege_level_name)(uint8_t pl);
    bool (*does_insn_modify_frame_pointer)(const char *disas);
} CpuOps;

typedef struct Cpu {
    Trace *trace;
    Callstack *cs;
    uint8_t privilege_level;
    GArray *traces; /* Trace *traces [] */
    GByteArray *buf;
    CpuOps ops;
    void *arch;
} Cpu;

typedef enum {
    AARCH64_EL0_SECURE,
    AARCH64_EL0_NONSECURE,
    AARCH64_EL0_REALM,
    AARCH64_EL1_SECURE,
    AARCH64_EL1_NONSECURE,
    AARCH64_EL1_REALM,
    AARCH64_EL2_SECURE,
    AARCH64_EL2_NONSECURE,
    AARCH64_EL2_REALM,
    AARCH64_EL3,
    AARCH64_PRIVILEGE_LEVEL_MAX,
} Aarch64PrivilegeLevel;

typedef struct {
    struct qemu_plugin_register *reg_fp;
    struct qemu_plugin_register *reg_cpsr;
    struct qemu_plugin_register *reg_scr_el3;
} Aarch64Cpu;

typedef enum {
    X64_RING0,
    X64_RING1,
    X64_RING2,
    X64_RING3,
    X64_REAL_MODE,
    X64_PRIVILEGE_LEVEL_MAX,
} X64PrivilegeLevel;

typedef struct {
    struct qemu_plugin_register *reg_rbp;
    struct qemu_plugin_register *reg_cs;
    struct qemu_plugin_register *reg_cr0;
} X64Cpu;

typedef struct {
    uint64_t timestamp;
    uint64_t data;
} UftraceEntry;

typedef enum {
    UFTRACE_ENTRY,
    UFTRACE_EXIT,
    UFTRACE_LOST,
    UFTRACE_EVENT,
} UftraceRecordType;

static struct qemu_plugin_scoreboard *score;
static bool trace_privilege_level;
static CpuOps arch_ops;

static uint64_t gettime_ns(void)
{
#ifdef _WIN32
    /*
     * On Windows, timespec_get is available only with UCRT, but not with
     * MinGW64 environment. Simplify by using only gettimeofday on this
     * platform. This may result in a precision loss.
     */
    struct timeval tv;
    gettimeofday(&tv, NULL);
    uint64_t now_ns = tv.tv_sec * NANOSECONDS_PER_SECOND + tv.tv_usec * 1000;
#else
    /* We need nanosecond precision for short lived functions. */
    struct timespec ts;
    timespec_get(&ts, TIME_UTC);
    uint64_t now_ns = ts.tv_sec * NANOSECONDS_PER_SECOND + ts.tv_nsec;
#endif
    return now_ns;
}

static void uftrace_write_map(bool system_emulation)
{
    const char *path = "./uftrace.data/sid-0.map";

    if (system_emulation && access(path, F_OK) == 0) {
        /* do not erase existing map in system emulation, as a custom one might
         * already have been generated by uftrace_symbols.py */
        return;
    }

    FILE *sid_map = fopen(path, "w");
    g_assert(sid_map);

    if (system_emulation) {
        fprintf(sid_map,
                "# map stack on highest address possible, to prevent uftrace\n"
                "# from considering any kernel address\n");
        fprintf(sid_map,
          "ffffffffffff-ffffffffffff rw-p 00000000 00:00 0 [stack]\n");
    } else {
        /* in user mode, copy /proc/self/maps instead */
        FILE *self_map = fopen("/proc/self/maps", "r");
        g_assert(self_map);
        for (;;) {
            int c = fgetc(self_map);
            if (c == EOF) {
                break;
            }
            fputc(c, sid_map);
        }
        fclose(self_map);
    }
    fclose(sid_map);
}

static void uftrace_write_task(const GArray *traces)
{
    FILE *task = fopen("./uftrace.data/task.txt", "w");
    g_assert(task);
    for (int i = 0; i < traces->len; ++i) {
        Trace *t = g_array_index(traces, Trace*, i);
        fprintf(task, "SESS timestamp=0.0 pid=%"PRIu32" sid=0 exename=\"%s\"\n",
                t->id, t->name->str);
        fprintf(task, "TASK timestamp=0.0 tid=%"PRIu32" pid=%"PRIu32"\n",
                t->id, t->id);
    }
    fclose(task);
}

static void uftrace_write_info(const GArray *traces)
{
    g_autoptr(GString) taskinfo_tids = g_string_new("taskinfo:tids=");
    for (int i = 0; i < traces->len; ++i) {
        Trace *t = g_array_index(traces, Trace*, i);
        const char *delim = i > 0 ? "," : "";
        g_string_append_printf(taskinfo_tids, "%s%"PRIu32, delim, t->id);
    }

    g_autoptr(GString) taskinfo_nr_tid = g_string_new("taskinfo:nr_tid=");
    g_string_append_printf(taskinfo_nr_tid, "%d", traces->len);

    FILE *info = fopen("./uftrace.data/info", "w");
    g_assert(info);
    /*
     * $ uftrace dump --debug
     * uftrace file header: magic         = 4674726163652100
     * uftrace file header: version       = 4
     * uftrace file header: header size   = 40
     * uftrace file header: endian        = 1 (little)
     * uftrace file header: class          = 2 (64 bit)
     * uftrace file header: features      = 0x1263 (PLTHOOK | ...
     * uftrace file header: info          = 0x7bff (EXE_NAME | ...
     *  <0000000000000000>: 46 74 72 61 63 65 21 00  04 00 00 00 28 00 01 02
     *  <0000000000000010>: 63 12 00 00 00 00 00 00  ff 7b 00 00 00 00 00 00
     *  <0000000000000020>: 00 04 00 00 00 00 00 00
     */
    const uint8_t header[] = {0x46, 0x74, 0x72, 0x61, 0x63, 0x65, 0x21, 0x00,
                              0x04, 0x00, 0x00, 0x00, 0x28, 0x00, 0x01, 0x02,
                              0x63, 0x12, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                              0xff, 0x7b, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                              0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};
    size_t wrote = fwrite(header, sizeof(header), 1, info);
    g_assert(wrote == 1);
    const char *info_data[] = {
        "exename:",
        "build_id:0000000000000000000000000000000000000000",
        "exit_status:",
        "cmdline:",
        "cpuinfo:lines=2",
        "cpuinfo:nr_cpus=",
        "cpuinfo:desc=",
        "meminfo:",
        "osinfo:lines=3",
        "osinfo:kernel=",
        "osinfo:hostname=",
        "osinfo:distro=",
        "taskinfo:lines=2",
        taskinfo_nr_tid->str,
        taskinfo_tids->str,
        "usageinfo:lines=6",
        "usageinfo:systime=",
        "usageinfo:usrtime=",
        "usageinfo:ctxsw=",
        "usageinfo:maxrss=",
        "usageinfo:pagefault=",
        "usageinfo:iops=",
        "loadinfo:",
        "record_date:",
        "elapsed_time:",
        "pattern_type:regex",
        "uftrace_version:",
        "utc_offset:",
        0};
    const char **info_data_it = info_data;
    while (*(info_data_it)) {
        fprintf(info, "%s\n", *info_data_it);
        ++info_data_it;
    }
    fclose(info);
}

static Callstack *callstack_new(void)
{
    Callstack *cs = g_new0(Callstack, 1);
    cs->s = g_array_new(false, false, sizeof(CallstackEntry));
    return cs;
}

static void callstack_free(Callstack *cs)
{
    g_array_free(cs->s, true);
    cs->s = NULL;
    g_free(cs);
}

static size_t callstack_depth(const Callstack *cs)
{
    return cs->s->len;
}

static size_t callstack_empty(const Callstack *cs)
{
    return callstack_depth(cs) == 0;
}

static void callstack_clear(Callstack *cs)
{
    g_array_set_size(cs->s, 0);
}

static const CallstackEntry *callstack_at(const Callstack *cs, size_t depth)
{
    g_assert(depth > 0);
    g_assert(depth <= callstack_depth(cs));
    return &g_array_index(cs->s, CallstackEntry, depth - 1);
}

static CallstackEntry callstack_top(const Callstack *cs)
{
    if (callstack_depth(cs) >= 1) {
        return *callstack_at(cs, callstack_depth(cs));
    }
    return (CallstackEntry){};
}

static CallstackEntry callstack_caller(const Callstack *cs)
{
    if (callstack_depth(cs) >= 2) {
        return *callstack_at(cs, callstack_depth(cs) - 1);
    }
    return (CallstackEntry){};
}

static void callstack_push(Callstack *cs, CallstackEntry e)
{
    g_array_append_val(cs->s, e);
}

static CallstackEntry callstack_pop(Callstack *cs)
{
    g_assert(!callstack_empty(cs));
    CallstackEntry e = callstack_top(cs);
    g_array_set_size(cs->s, callstack_depth(cs) - 1);
    return e;
}

static Trace *trace_new(uint32_t id, GString *name)
{
    Trace *t = g_new0(Trace, 1);
    t->t = g_array_new(false, false, sizeof(UftraceEntry));
    t->path = g_string_new(NULL);
    g_string_append_printf(t->path, "./uftrace.data/%"PRIu32".dat", id);
    t->name = g_string_new(name->str);
    t->id = id;
    return t;
}

static void trace_free(Trace *t)
{
    g_assert(t->t->len == 0);
    g_array_free(t->t, true);
    t->t = NULL;
    g_string_free(t->path, true);
    t->path = NULL;
    g_string_free(t->name, true);
    t->name = NULL;
    g_free(t);
}

static void trace_flush(Trace *t, bool append)
{
    int create_dir = g_mkdir_with_parents("./uftrace.data",
                                          S_IRWXU | S_IRWXG | S_IRWXO);
    g_assert(create_dir == 0);
    FILE *dat = fopen(t->path->str, append ? "a" : "w");
    g_assert(dat);
    GArray *data = t->t;
    if (data->len) {
        size_t wrote = fwrite(data->data, sizeof(UftraceEntry), data->len, dat);
        g_assert(wrote == data->len);
    }
    fclose(dat);
    g_array_set_size(data, 0);
}

static void trace_add_entry(Trace *t, uint64_t timestamp, uint64_t pc,
                            size_t depth, UftraceRecordType type)
{
    /* https://github.com/namhyung/uftrace/blob/v0.18/libmcount/record.c#L909 */
    const uint64_t record_magic = 0x5;
    uint64_t data = type | (record_magic << 3);
    data += depth << 6;
    data += pc << 16;
    UftraceEntry e = {.timestamp = timestamp, .data = data};
    g_array_append_val(t->t, e);
    if (t->t->len * sizeof(UftraceEntry) > TRACE_FLUSH_SIZE) {
        trace_flush(t, true);
    }
}

static void trace_enter_function(Trace *t, uint64_t timestamp,
                                 uint64_t pc, size_t depth)
{
    trace_add_entry(t, timestamp, pc, depth, UFTRACE_ENTRY);
}

static void trace_exit_function(Trace *t, uint64_t timestamp,
                                uint64_t pc, size_t depth)
{
    trace_add_entry(t, timestamp, pc, depth, UFTRACE_EXIT);
}

static void trace_enter_stack(Trace *t, Callstack *cs, uint64_t timestamp)
{
    for (size_t depth = 1; depth <= callstack_depth(cs); ++depth) {
        trace_enter_function(t, timestamp, callstack_at(cs, depth)->pc, depth);
    }
}

static void trace_exit_stack(Trace *t, Callstack *cs, uint64_t timestamp)
{
    for (size_t depth = callstack_depth(cs); depth > 0; --depth) {
        trace_exit_function(t, timestamp, callstack_at(cs, depth)->pc, depth);
    }
}

static uint64_t cpu_read_register64(Cpu *cpu, struct qemu_plugin_register *reg)
{
    GByteArray *buf = cpu->buf;
    g_byte_array_set_size(buf, 0);
    size_t sz = qemu_plugin_read_register(reg, buf);
    g_assert(sz == 8);
    g_assert(buf->len == 8);
    return *((uint64_t *) buf->data);
}

static uint32_t cpu_read_register32(Cpu *cpu, struct qemu_plugin_register *reg)
{
    GByteArray *buf = cpu->buf;
    g_byte_array_set_size(buf, 0);
    size_t sz = qemu_plugin_read_register(reg, buf);
    g_assert(sz == 4);
    g_assert(buf->len == 4);
    return *((uint32_t *) buf->data);
}

static uint64_t cpu_read_memory64(Cpu *cpu, uint64_t addr)
{
    g_assert(addr);
    GByteArray *buf = cpu->buf;
    g_byte_array_set_size(buf, 0);
    bool read = qemu_plugin_read_memory_vaddr(addr, buf, 8);
    if (!read) {
        return 0;
    }
    g_assert(buf->len == 8);
    return *((uint64_t *) buf->data);
}

static void cpu_unwind_stack(Cpu *cpu, uint64_t frame_pointer, uint64_t pc)
{
    g_assert(callstack_empty(cpu->cs));

    #define UNWIND_STACK_MAX_DEPTH 1024
    CallstackEntry unwind[UNWIND_STACK_MAX_DEPTH];
    size_t depth = 0;
    do {
        /* check we don't have an infinite stack */
        for (size_t i = 0; i < depth; ++i) {
            if (frame_pointer == unwind[i].frame_pointer) {
                break;
            }
        }
        CallstackEntry e = {.frame_pointer = frame_pointer, .pc = pc};
        unwind[depth] = e;
        depth++;
        if (frame_pointer) {
            frame_pointer = cpu_read_memory64(cpu, frame_pointer);
        }
        pc = cpu_read_memory64(cpu, frame_pointer + 8); /* read previous lr */
    } while (frame_pointer && pc && depth < UNWIND_STACK_MAX_DEPTH);
    #undef UNWIND_STACK_MAX_DEPTH

    /* push it from bottom to top */
    while (depth) {
        callstack_push(cpu->cs, unwind[depth - 1]);
        --depth;
    }
}

static struct qemu_plugin_register *plugin_find_register(const char *name)
{
    g_autoptr(GArray) regs = qemu_plugin_get_registers();
    for (int i = 0; i < regs->len; ++i) {
        qemu_plugin_reg_descriptor *reg;
        reg = &g_array_index(regs, qemu_plugin_reg_descriptor, i);
        if (!strcmp(reg->name, name)) {
            return reg->handle;
        }
    }
    return NULL;
}

static uint8_t aarch64_num_privilege_levels(void)
{
    return AARCH64_PRIVILEGE_LEVEL_MAX;
}

static const char *aarch64_get_privilege_level_name(uint8_t pl)
{
    switch (pl) {
    case AARCH64_EL0_SECURE: return "S-EL0";
    case AARCH64_EL0_NONSECURE: return "NS-EL0";
    case AARCH64_EL0_REALM: return "R-EL0";
    case AARCH64_EL1_SECURE: return "S-EL1";
    case AARCH64_EL1_NONSECURE: return "NS-EL1";
    case AARCH64_EL1_REALM: return "R-EL1";
    case AARCH64_EL2_SECURE: return "S-EL2";
    case AARCH64_EL2_NONSECURE: return "NS-EL2";
    case AARCH64_EL2_REALM: return "R-EL2";
    case AARCH64_EL3: return "EL3";
    default:
        g_assert_not_reached();
    }
}

static uint8_t aarch64_get_privilege_level(Cpu *cpu_)
{
    Aarch64Cpu *cpu = cpu_->arch;
    /*
     * QEMU gdbstub does not provide access to CurrentEL,
     * so we use CPSR instead.
     */
    uint8_t el = cpu_read_register32(cpu_, cpu->reg_cpsr) >> 2 & 0b11;

    if (el == 3) {
        return AARCH64_EL3;
    }

    uint8_t ss = AARCH64_EL0_SECURE;
    if (!cpu->reg_scr_el3) {
        ss = AARCH64_EL0_NONSECURE;
    }
    uint64_t scr_el3 = cpu_read_register64(cpu_, cpu->reg_scr_el3);
    uint64_t ns = (scr_el3 >> 0) & 0b1;
    uint64_t nse = (scr_el3 >> 62) & 0b1;
    switch (nse << 1 | ns) {
    case 0b00:
        ss = AARCH64_EL0_SECURE;
        break;
    case 0b01:
        ss = AARCH64_EL0_NONSECURE;
        break;
    case 0b11:
        ss = AARCH64_EL0_REALM;
        break;
    default:
        g_assert_not_reached();
    }

    const uint8_t num_ss = 3;
    Aarch64PrivilegeLevel pl = el * num_ss + ss;
    return pl;
}

static uint64_t aarch64_get_frame_pointer(Cpu *cpu_)
{
    Aarch64Cpu *cpu = cpu_->arch;
    return cpu_read_register64(cpu_, cpu->reg_fp);
}

static void aarch64_init(Cpu *cpu_)
{
    Aarch64Cpu *cpu = g_new0(Aarch64Cpu, 1);
    cpu_->arch = cpu;
    cpu->reg_fp = plugin_find_register("x29");
    if (!cpu->reg_fp) {
        fprintf(stderr, "uftrace plugin: frame pointer register (x29) is not "
                        "available. Please use an AArch64 cpu (or -cpu max).\n");
        g_abort();
    }
    cpu->reg_cpsr = plugin_find_register("cpsr");
    g_assert(cpu->reg_cpsr);
    cpu->reg_scr_el3 = plugin_find_register("SCR_EL3");
    /* scr_el3 is optional */
}

static void aarch64_end(Cpu *cpu)
{
    g_free(cpu->arch);
}

static bool aarch64_does_insn_modify_frame_pointer(const char *disas)
{
    /*
     * Check if current instruction concerns fp register "x29".
     * We add a prefix space to make sure we don't match addresses dump
     * in disassembly.
     */
    return strstr(disas, " x29");
}

static CpuOps aarch64_ops = {
    .init = aarch64_init,
    .end = aarch64_end,
    .get_frame_pointer = aarch64_get_frame_pointer,
    .get_privilege_level = aarch64_get_privilege_level,
    .num_privilege_levels = aarch64_num_privilege_levels,
    .get_privilege_level_name = aarch64_get_privilege_level_name,
    .does_insn_modify_frame_pointer = aarch64_does_insn_modify_frame_pointer,
};

static uint8_t x64_num_privilege_levels(void)
{
    return X64_PRIVILEGE_LEVEL_MAX;
}

static const char *x64_get_privilege_level_name(uint8_t pl)
{
    switch (pl) {
    case X64_RING0: return "Ring0";
    case X64_RING1: return "Ring1";
    case X64_RING2: return "Ring2";
    case X64_RING3: return "Ring3";
    case X64_REAL_MODE: return "RealMode";
    default:
        g_assert_not_reached();
    }
}

static uint8_t x64_get_privilege_level(Cpu *cpu_)
{
    X64Cpu *cpu = cpu_->arch;
    uint64_t cr0 = cpu_read_register64(cpu_, cpu->reg_cr0);
    uint64_t protected_mode = (cr0 >> 0) & 0b1;
    if (!protected_mode) {
        return X64_REAL_MODE;
    }
    uint32_t cs = cpu_read_register32(cpu_, cpu->reg_cs);
    uint32_t ring_level = (cs >> 0) & 0b11;
    return ring_level;
}

static uint64_t x64_get_frame_pointer(Cpu *cpu_)
{
    X64Cpu *cpu = cpu_->arch;
    return cpu_read_register64(cpu_, cpu->reg_rbp);
}

static void x64_init(Cpu *cpu_)
{
    X64Cpu *cpu = g_new0(X64Cpu, 1);
    cpu_->arch = cpu;
    cpu->reg_rbp = plugin_find_register("rbp");
    g_assert(cpu->reg_rbp);
    cpu->reg_cs = plugin_find_register("cs");
    g_assert(cpu->reg_cs);
    cpu->reg_cr0 = plugin_find_register("cr0");
    g_assert(cpu->reg_cr0);
}

static void x64_end(Cpu *cpu)
{
    g_free(cpu->arch);
}

static bool x64_does_insn_modify_frame_pointer(const char *disas)
{
    return strstr(disas, "rbp");
}

static CpuOps x64_ops = {
    .init = x64_init,
    .end = x64_end,
    .get_frame_pointer = x64_get_frame_pointer,
    .get_privilege_level = x64_get_privilege_level,
    .num_privilege_levels = x64_num_privilege_levels,
    .get_privilege_level_name = x64_get_privilege_level_name,
    .does_insn_modify_frame_pointer = x64_does_insn_modify_frame_pointer,
};

static void track_privilege_change(unsigned int cpu_index, void *udata)
{
    Cpu *cpu = qemu_plugin_scoreboard_find(score, cpu_index);
    uint8_t new_pl = cpu->ops.get_privilege_level(cpu);

    if (new_pl == cpu->privilege_level) {
        return;
    }

    uint64_t pc = (uintptr_t) udata;
    uint64_t timestamp = gettime_ns();

    trace_exit_stack(cpu->trace, cpu->cs, timestamp);
    callstack_clear(cpu->cs);

    cpu->privilege_level = new_pl;
    cpu->trace = g_array_index(cpu->traces, Trace*, new_pl);

    cpu_unwind_stack(cpu, cpu->ops.get_frame_pointer(cpu), pc);
    trace_enter_stack(cpu->trace, cpu->cs, timestamp);
}

static void track_callstack(unsigned int cpu_index, void *udata)
{
    uint64_t pc = (uintptr_t) udata;
    Cpu *cpu = qemu_plugin_scoreboard_find(score, cpu_index);
    uint64_t timestamp = gettime_ns();
    Callstack *cs = cpu->cs;
    Trace *t = cpu->trace;

    uint64_t fp = cpu->ops.get_frame_pointer(cpu);
    if (!fp && callstack_empty(cs)) {
        /*
         * We simply push current pc. Note that we won't detect symbol change as
         * long as a proper call does not happen.
         */
        callstack_push(cs, (CallstackEntry){.frame_pointer = fp, .pc = pc});
        trace_enter_function(t, timestamp, pc, callstack_depth(cs));
        return;
    }

    CallstackEntry top = callstack_top(cs);
    if (fp == top.frame_pointer) {
        /* same function */
        return;
    }

    CallstackEntry caller = callstack_caller(cs);
    if (fp == caller.frame_pointer) {
        /* return */
        CallstackEntry e = callstack_pop(cs);
        trace_exit_function(t, timestamp, e.pc, callstack_depth(cs));
        return;
    }

    uint64_t caller_fp = fp ? cpu_read_memory64(cpu, fp) : 0;
    if (caller_fp == top.frame_pointer) {
        /* call */
        callstack_push(cs, (CallstackEntry){.frame_pointer = fp, .pc = pc});
        trace_enter_function(t, timestamp, pc, callstack_depth(cs));
        return;
    }

    /* discontinuity, exit current stack and unwind new one */
    trace_exit_stack(t, cs, timestamp);
    callstack_clear(cs);

    cpu_unwind_stack(cpu, fp, pc);
    trace_enter_stack(t, cs, timestamp);
}

static void vcpu_tb_trans(qemu_plugin_id_t id, struct qemu_plugin_tb *tb)
{
    size_t n_insns = qemu_plugin_tb_n_insns(tb);
    uintptr_t tb_pc = qemu_plugin_tb_vaddr(tb);

    if (trace_privilege_level) {
        qemu_plugin_register_vcpu_tb_exec_cb(tb, track_privilege_change,
                                             QEMU_PLUGIN_CB_R_REGS,
                                             (void *) tb_pc);
    }

    /*
     * Callbacks and inline instrumentation are inserted before an instruction.
     * Thus, to see instruction effect, we need to wait for next one.
     * Potentially, the last instruction of a block could modify the frame
     * pointer. Thus, we need to always instrument first instruction in a tb.
     */
    bool instrument_insn = true;
    for (size_t i = 0; i < n_insns; i++) {
        struct qemu_plugin_insn *insn = qemu_plugin_tb_get_insn(tb, i);

        if (instrument_insn) {
            uintptr_t pc = qemu_plugin_insn_vaddr(insn);
            qemu_plugin_register_vcpu_insn_exec_cb(insn, track_callstack,
                                                   QEMU_PLUGIN_CB_R_REGS,
                                                   (void *) pc);
            instrument_insn = false;
        }

        char *disas = qemu_plugin_insn_disas(insn);
        if (arch_ops.does_insn_modify_frame_pointer(disas)) {
            instrument_insn = true;
        }
    }
}

static void vcpu_init(qemu_plugin_id_t id, unsigned int vcpu_index)
{
    Cpu *cpu = qemu_plugin_scoreboard_find(score, vcpu_index);
    cpu->ops = arch_ops;

    cpu->ops.init(cpu);
    cpu->buf = g_byte_array_new();
    cpu->traces = g_array_new(0, 0, sizeof(Trace *));

    g_assert(vcpu_index < UINT32_MAX / TRACE_ID_SCALE);
    g_assert(cpu->ops.num_privilege_levels() < TRACE_ID_SCALE);
    /* trace_id is: cpu_number * TRACE_ID_SCALE + privilege_level */
    uint32_t trace_id = (vcpu_index + 1) * TRACE_ID_SCALE;

    if (trace_privilege_level) {
        for (uint8_t pl = 0; pl < cpu->ops.num_privilege_levels(); ++pl) {
            g_autoptr(GString) trace_name = g_string_new(NULL);
            g_string_append_printf(trace_name, "cpu%u %s", vcpu_index,
                                   cpu->ops.get_privilege_level_name(pl));
            Trace *t = trace_new(trace_id + pl, trace_name);
            g_array_append_val(cpu->traces, t);
        }
    } else {
        g_autoptr(GString) trace_name = g_string_new(NULL);
        g_string_append_printf(trace_name, "cpu%u", vcpu_index);
        Trace *t = trace_new(trace_id, trace_name);
        g_array_append_val(cpu->traces, t);
    }

    for (size_t i = 0; i < cpu->traces->len; ++i) {
        /* create/truncate trace files */
        Trace *t = g_array_index(cpu->traces, Trace*, i);
        trace_flush(t, false);
    }

    cpu->cs = callstack_new();
    cpu->trace = g_array_index(cpu->traces, Trace*, cpu->privilege_level);
}

static void vcpu_end(unsigned int vcpu_index)
{
    Cpu *cpu = qemu_plugin_scoreboard_find(score, vcpu_index);
    g_byte_array_free(cpu->buf, true);

    for (size_t i = 0; i < cpu->traces->len; ++i) {
        Trace *t = g_array_index(cpu->traces, Trace*, i);
        trace_free(t);
    }

    g_array_free(cpu->traces, true);
    callstack_free(cpu->cs);
    memset(cpu, 0, sizeof(Cpu));
}

static void at_exit(qemu_plugin_id_t id, void *data)
{
    bool system_emulation = (bool) data;
    g_autoptr(GArray) traces = g_array_new(0, 0, sizeof(Trace *));

    for (size_t i = 0; i < qemu_plugin_num_vcpus(); ++i) {
        Cpu *cpu = qemu_plugin_scoreboard_find(score, i);
        for (size_t j = 0; j < cpu->traces->len; ++j) {
            Trace *t = g_array_index(cpu->traces, Trace*, j);
            trace_flush(t, true);
            g_array_append_val(traces, t);
        }
    }

    uftrace_write_map(system_emulation);
    uftrace_write_info(traces);
    uftrace_write_task(traces);

    for (size_t i = 0; i < qemu_plugin_num_vcpus(); ++i) {
        vcpu_end(i);
    }

    qemu_plugin_scoreboard_free(score);
}

QEMU_PLUGIN_EXPORT int qemu_plugin_install(qemu_plugin_id_t id,
                                           const qemu_info_t *info,
                                           int argc, char **argv)
{
    for (int i = 0; i < argc; i++) {
        char *opt = argv[i];
        g_auto(GStrv) tokens = g_strsplit(opt, "=", 2);
        if (g_strcmp0(tokens[0], "trace-privilege-level") == 0) {
            if (!qemu_plugin_bool_parse(tokens[0], tokens[1],
                                        &trace_privilege_level)) {
                fprintf(stderr, "boolean argument parsing failed: %s\n", opt);
                return -1;
            }
        } else {
            fprintf(stderr, "option parsing failed: %s\n", opt);
            return -1;
        }
    }

    if (!strcmp(info->target_name, "aarch64")) {
        arch_ops = aarch64_ops;
    } else if (!strcmp(info->target_name, "x86_64")) {
        arch_ops = x64_ops;
    } else {
        fprintf(stderr, "plugin uftrace: %s target is not supported\n",
                info->target_name);
        return 1;
    }

    score = qemu_plugin_scoreboard_new(sizeof(Cpu));
    qemu_plugin_register_vcpu_init_cb(id, vcpu_init);
    qemu_plugin_register_atexit_cb(id, at_exit, (void *) info->system_emulation);
    qemu_plugin_register_vcpu_tb_trans_cb(id, vcpu_tb_trans);

    return 0;
}