#include <sys/types.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <inttypes.h>
#include <byteswap.h>
#include <sys/stat.h>
#include <sys/mman.h>

#include "util.h"
#include "event.h"
#include "debug.h"
#include "evlist.h"
#include "symbol.h"
#include "strlist.h"
#include <elf.h>

#include "session.h"
#include "jit.h"
#include "jitdump.h"
#include "genelf.h"
#include "../builtin.h"

struct jit_buf_desc {
        struct perf_data_file *output;
        struct perf_session *session;
        struct machine *machine;
        union jr_entry   *entry;
        void             *buf;
        uint64_t         sample_type;
        size_t           bufsize;
        FILE             *in;
        bool             needs_bswap; /* handles cross-endianess */
        void             *debug_data;
        size_t           nr_debug_entries;
        uint32_t         code_load_count;
        u64              bytes_written;
        struct rb_root   code_root;
        char             dir[PATH_MAX];
};

struct debug_line_info {
        unsigned long vma;
        unsigned int lineno;
        /* The filename format is unspecified, absolute path, relative etc. */
        char const filename[0];
};

struct jit_tool {
        struct perf_tool tool;
        struct perf_data_file   output;
        struct perf_data_file   input;
        u64 bytes_written;
};

#define hmax(a, b) ((a) > (b) ? (a) : (b))
#define get_jit_tool(t) (container_of(tool, struct jit_tool, tool))

static int
jit_emit_elf(char *filename,
             const char *sym,
             uint64_t code_addr,
             const void *code,
             int csize,
             void *debug,
             int nr_debug_entries)
{
        int ret, fd;

        if (verbose > 0)
                fprintf(stderr, "write ELF image %s\n", filename);

        fd = open(filename, O_CREAT|O_TRUNC|O_WRONLY, 0644);
        if (fd == -1) {
                pr_warning("cannot create jit ELF %s: %s\n", filename, strerror(errno));
                return -1;
        }

        ret = jit_write_elf(fd, code_addr, sym, (const void *)code, csize, debug, nr_debug_entries);

        close(fd);

        if (ret)
                unlink(filename);

        return ret;
}

static void
jit_close(struct jit_buf_desc *jd)
{
        if (!(jd && jd->in))
                return;
        funlockfile(jd->in);
        fclose(jd->in);
        jd->in = NULL;
}

static int
jit_validate_events(struct perf_session *session)
{
        struct perf_evsel *evsel;

        /*
         * check that all events use CLOCK_MONOTONIC
         */
        evlist__for_each(session->evlist, evsel) {
                if (evsel->attr.use_clockid == 0 || evsel->attr.clockid != CLOCK_MONOTONIC)
                        return -1;
        }
        return 0;
}

static int
jit_open(struct jit_buf_desc *jd, const char *name)
{
        struct jitheader header;
        struct jr_prefix *prefix;
        ssize_t bs, bsz = 0;
        void *n, *buf = NULL;
        int ret, retval = -1;

        jd->in = fopen(name, "r");
        if (!jd->in)
                return -1;

        bsz = hmax(sizeof(header), sizeof(*prefix));

        buf = malloc(bsz);
        if (!buf)
                goto error;

        /*
         * protect from writer modifying the file while we are reading it
         */
        flockfile(jd->in);

        ret = fread(buf, sizeof(header), 1, jd->in);
        if (ret != 1)
                goto error;

        memcpy(&header, buf, sizeof(header));

        if (header.magic != JITHEADER_MAGIC) {
                if (header.magic != JITHEADER_MAGIC_SW)
                        goto error;
                jd->needs_bswap = true;
        }

        if (jd->needs_bswap) {
                header.version    = bswap_32(header.version);
                header.total_size = bswap_32(header.total_size);
                header.pid        = bswap_32(header.pid);
                header.elf_mach   = bswap_32(header.elf_mach);
                header.timestamp  = bswap_64(header.timestamp);
                header.flags      = bswap_64(header.flags);
        }

        if (verbose > 2)
                pr_debug("version=%u\nhdr.size=%u\nts=0x%llx\npid=%d\nelf_mach=%d\n",
                        header.version,
                        header.total_size,
                        (unsigned long long)header.timestamp,
                        header.pid,
                        header.elf_mach);

        if (header.flags & JITDUMP_FLAGS_RESERVED) {
                pr_err("jitdump file contains invalid or unsupported flags 0x%llx\n",
                       (unsigned long long)header.flags & JITDUMP_FLAGS_RESERVED);
                goto error;
        }

        /*
         * validate event is using the correct clockid
         */
        if (jit_validate_events(jd->session)) {
                pr_err("error, jitted code must be sampled with perf record -k 1\n");
                goto error;
        }

        bs = header.total_size - sizeof(header);

        if (bs > bsz) {
                n = realloc(buf, bs);
                if (!n)
                        goto error;
                bsz = bs;
                buf = n;
                /* read extra we do not know about */
                ret = fread(buf, bs - bsz, 1, jd->in);
                if (ret != 1)
                        goto error;
        }
        /*
         * keep dirname for generating files and mmap records
         */
        strcpy(jd->dir, name);
        dirname(jd->dir);

        return 0;
error:
        funlockfile(jd->in);
        fclose(jd->in);
        return retval;
}

static union jr_entry *
jit_get_next_entry(struct jit_buf_desc *jd)
{
        struct jr_prefix *prefix;
        union jr_entry *jr;
        void *addr;
        size_t bs, size;
        int id, ret;

        if (!(jd && jd->in))
                return NULL;

        if (jd->buf == NULL) {
                size_t sz = getpagesize();
                if (sz < sizeof(*prefix))
                        sz = sizeof(*prefix);

                jd->buf = malloc(sz);
                if (jd->buf == NULL)
                        return NULL;

                jd->bufsize = sz;
        }

        prefix = jd->buf;

        /*
         * file is still locked at this point
         */
        ret = fread(prefix, sizeof(*prefix), 1, jd->in);
        if (ret  != 1)
                return NULL;

        if (jd->needs_bswap) {
                prefix->id         = bswap_32(prefix->id);
                prefix->total_size = bswap_32(prefix->total_size);
                prefix->timestamp  = bswap_64(prefix->timestamp);
        }
        id   = prefix->id;
        size = prefix->total_size;

        bs = (size_t)size;
        if (bs < sizeof(*prefix))
                return NULL;

        if (id >= JIT_CODE_MAX) {
                pr_warning("next_entry: unknown prefix %d, skipping\n", id);
                return NULL;
        }
        if (bs > jd->bufsize) {
                void *n;
                n = realloc(jd->buf, bs);
                if (!n)
                        return NULL;
                jd->buf = n;
                jd->bufsize = bs;
        }

        addr = ((void *)jd->buf) + sizeof(*prefix);

        ret = fread(addr, bs - sizeof(*prefix), 1, jd->in);
        if (ret != 1)
                return NULL;

        jr = (union jr_entry *)jd->buf;

        switch(id) {
        case JIT_CODE_DEBUG_INFO:
                if (jd->needs_bswap) {
                        uint64_t n;
                        jr->info.code_addr = bswap_64(jr->info.code_addr);
                        jr->info.nr_entry  = bswap_64(jr->info.nr_entry);
                        for (n = 0 ; n < jr->info.nr_entry; n++) {
                                jr->info.entries[n].addr    = bswap_64(jr->info.entries[n].addr);
                                jr->info.entries[n].lineno  = bswap_32(jr->info.entries[n].lineno);
                                jr->info.entries[n].discrim = bswap_32(jr->info.entries[n].discrim);
                        }
                }
                break;
        case JIT_CODE_CLOSE:
                break;
        case JIT_CODE_LOAD:
                if (jd->needs_bswap) {
                        jr->load.pid       = bswap_32(jr->load.pid);
                        jr->load.tid       = bswap_32(jr->load.tid);
                        jr->load.vma       = bswap_64(jr->load.vma);
                        jr->load.code_addr = bswap_64(jr->load.code_addr);
                        jr->load.code_size = bswap_64(jr->load.code_size);
                        jr->load.code_index= bswap_64(jr->load.code_index);
                }
                jd->code_load_count++;
                break;
        case JIT_CODE_MOVE:
                if (jd->needs_bswap) {
                        jr->move.pid           = bswap_32(jr->move.pid);
                        jr->move.tid           = bswap_32(jr->move.tid);
                        jr->move.vma           = bswap_64(jr->move.vma);
                        jr->move.old_code_addr = bswap_64(jr->move.old_code_addr);
                        jr->move.new_code_addr = bswap_64(jr->move.new_code_addr);
                        jr->move.code_size     = bswap_64(jr->move.code_size);
                        jr->move.code_index    = bswap_64(jr->move.code_index);
                }
                break;
        case JIT_CODE_MAX:
        default:
                return NULL;
        }
        return jr;
}

static int
jit_inject_event(struct jit_buf_desc *jd, union perf_event *event)
{
        ssize_t size;

        size = perf_data_file__write(jd->output, event, event->header.size);
        if (size < 0)
                return -1;

        jd->bytes_written += size;
        return 0;
}

static int jit_repipe_code_load(struct jit_buf_desc *jd, union jr_entry *jr)
{
        struct perf_sample sample;
        union perf_event *event;
        struct perf_tool *tool = jd->session->tool;
        uint64_t code, addr;
        uintptr_t uaddr;
        char *filename;
        struct stat st;
        size_t size;
        u16 idr_size;
        const char *sym;
        uint32_t count;
        int ret, csize;
        pid_t pid, tid;
        struct {
                u32 pid, tid;
                u64 time;
        } *id;

        pid   = jr->load.pid;
        tid   = jr->load.tid;
        csize = jr->load.code_size;
        addr  = jr->load.code_addr;
        sym   = (void *)((unsigned long)jr + sizeof(jr->load));
        code  = (unsigned long)jr + jr->load.p.total_size - csize;
        count = jr->load.code_index;
        idr_size = jd->machine->id_hdr_size;

        event = calloc(1, sizeof(*event) + idr_size);
        if (!event)
                return -1;

        filename = event->mmap2.filename;
        size = snprintf(filename, PATH_MAX, "%s/jitted-%d-%u.so",
                        jd->dir,
                        pid,
                        count);

        size++; /* for \0 */

        size = PERF_ALIGN(size, sizeof(u64));
        uaddr = (uintptr_t)code;
        ret = jit_emit_elf(filename, sym, addr, (const void *)uaddr, csize, jd->debug_data, jd->nr_debug_entries);

        if (jd->debug_data && jd->nr_debug_entries) {
                free(jd->debug_data);
                jd->debug_data = NULL;
                jd->nr_debug_entries = 0;
        }

        if (ret) {
                free(event);
                return -1;
        }
        if (stat(filename, &st))
                memset(&st, 0, sizeof(stat));

        event->mmap2.header.type = PERF_RECORD_MMAP2;
        event->mmap2.header.misc = PERF_RECORD_MISC_USER;
        event->mmap2.header.size = (sizeof(event->mmap2) -
                        (sizeof(event->mmap2.filename) - size) + idr_size);

        event->mmap2.pgoff = GEN_ELF_TEXT_OFFSET;
        event->mmap2.start = addr;
        event->mmap2.len   = csize;
        event->mmap2.pid   = pid;
        event->mmap2.tid   = tid;
        event->mmap2.ino   = st.st_ino;
        event->mmap2.maj   = major(st.st_dev);
        event->mmap2.min   = minor(st.st_dev);
        event->mmap2.prot  = st.st_mode;
        event->mmap2.flags = MAP_SHARED;
        event->mmap2.ino_generation = 1;

        id = (void *)((unsigned long)event + event->mmap.header.size - idr_size);
        if (jd->sample_type & PERF_SAMPLE_TID) {
                id->pid  = pid;
                id->tid  = tid;
        }
        if (jd->sample_type & PERF_SAMPLE_TIME)
                id->time = jr->load.p.timestamp;

        /*
         * create pseudo sample to induce dso hit increment
         * use first address as sample address
         */
        memset(&sample, 0, sizeof(sample));
        sample.cpumode = PERF_RECORD_MISC_USER;
        sample.pid  = pid;
        sample.tid  = tid;
        sample.time = id->time;
        sample.ip   = addr;

        ret = perf_event__process_mmap2(tool, event, &sample, jd->machine);
        if (ret)
                return ret;

        ret = jit_inject_event(jd, event);
        /*
         * mark dso as use to generate buildid in the header
         */
        if (!ret)
                build_id__mark_dso_hit(tool, event, &sample, NULL, jd->machine);

        return ret;
}

static int jit_repipe_code_move(struct jit_buf_desc *jd, union jr_entry *jr)
{
        struct perf_sample sample;
        union perf_event *event;
        struct perf_tool *tool = jd->session->tool;
        char *filename;
        size_t size;
        struct stat st;
        u16 idr_size;
        int ret;
        pid_t pid, tid;
        struct {
                u32 pid, tid;
                u64 time;
        } *id;

        pid = jr->move.pid;
        tid =  jr->move.tid;
        idr_size = jd->machine->id_hdr_size;

        /*
         * +16 to account for sample_id_all (hack)
         */
        event = calloc(1, sizeof(*event) + 16);
        if (!event)
                return -1;

        filename = event->mmap2.filename;
        size = snprintf(filename, PATH_MAX, "%s/jitted-%d-%"PRIu64,
                 jd->dir,
                 pid,
                 jr->move.code_index);

        size++; /* for \0 */

        if (stat(filename, &st))
                memset(&st, 0, sizeof(stat));

        size = PERF_ALIGN(size, sizeof(u64));

        event->mmap2.header.type = PERF_RECORD_MMAP2;
        event->mmap2.header.misc = PERF_RECORD_MISC_USER;
        event->mmap2.header.size = (sizeof(event->mmap2) -
                        (sizeof(event->mmap2.filename) - size) + idr_size);
        event->mmap2.pgoff = GEN_ELF_TEXT_OFFSET;
        event->mmap2.start = jr->move.new_code_addr;
        event->mmap2.len   = jr->move.code_size;
        event->mmap2.pid   = pid;
        event->mmap2.tid   = tid;
        event->mmap2.ino   = st.st_ino;
        event->mmap2.maj   = major(st.st_dev);
        event->mmap2.min   = minor(st.st_dev);
        event->mmap2.prot  = st.st_mode;
        event->mmap2.flags = MAP_SHARED;
        event->mmap2.ino_generation = 1;

        id = (void *)((unsigned long)event + event->mmap.header.size - idr_size);
        if (jd->sample_type & PERF_SAMPLE_TID) {
                id->pid  = pid;
                id->tid  = tid;
        }
        if (jd->sample_type & PERF_SAMPLE_TIME)
                id->time = jr->load.p.timestamp;

        /*
         * create pseudo sample to induce dso hit increment
         * use first address as sample address
         */
        memset(&sample, 0, sizeof(sample));
        sample.cpumode = PERF_RECORD_MISC_USER;
        sample.pid  = pid;
        sample.tid  = tid;
        sample.time = id->time;
        sample.ip   = jr->move.new_code_addr;

        ret = perf_event__process_mmap2(tool, event, &sample, jd->machine);
        if (ret)
                return ret;

        ret = jit_inject_event(jd, event);
        if (!ret)
                build_id__mark_dso_hit(tool, event, &sample, NULL, jd->machine);

        return ret;
}

static int jit_repipe_debug_info(struct jit_buf_desc *jd, union jr_entry *jr)
{
        void *data;
        size_t sz;

        if (!(jd && jr))
                return -1;

        sz  = jr->prefix.total_size - sizeof(jr->info);
        data = malloc(sz);
        if (!data)
                return -1;

        memcpy(data, &jr->info.entries, sz);

        jd->debug_data       = data;

        /*
         * we must use nr_entry instead of size here because
         * we cannot distinguish actual entry from padding otherwise
         */
        jd->nr_debug_entries = jr->info.nr_entry;

        return 0;
}

static int
jit_process_dump(struct jit_buf_desc *jd)
{
        union jr_entry *jr;
        int ret;

        while ((jr = jit_get_next_entry(jd))) {
                switch(jr->prefix.id) {
                case JIT_CODE_LOAD:
                        ret = jit_repipe_code_load(jd, jr);
                        break;
                case JIT_CODE_MOVE:
                        ret = jit_repipe_code_move(jd, jr);
                        break;
                case JIT_CODE_DEBUG_INFO:
                        ret = jit_repipe_debug_info(jd, jr);
                        break;
                default:
                        ret = 0;
                        continue;
                }
        }
        return ret;
}

static int
jit_inject(struct jit_buf_desc *jd, char *path)
{
        int ret;

        if (verbose > 0)
                fprintf(stderr, "injecting: %s\n", path);

        ret = jit_open(jd, path);
        if (ret)
                return -1;

        ret = jit_process_dump(jd);

        jit_close(jd);

        if (verbose > 0)
                fprintf(stderr, "injected: %s (%d)\n", path, ret);

        return 0;
}

/*
 * File must be with pattern .../jit-XXXX.dump
 * where XXXX is the PID of the process which did the mmap()
 * as captured in the RECORD_MMAP record
 */
static int
jit_detect(char *mmap_name, pid_t pid)
 {
        char *p;
        char *end = NULL;
        pid_t pid2;

        if (verbose > 2)
                fprintf(stderr, "jit marker trying : %s\n", mmap_name);
        /*
         * get file name
         */
        p = strrchr(mmap_name, '/');
        if (!p)
                return -1;

        /*
         * match prefix
         */
        if (strncmp(p, "/jit-", 5))
                return -1;

        /*
         * skip prefix
         */
        p += 5;

        /*
         * must be followed by a pid
         */
        if (!isdigit(*p))
                return -1;

        pid2 = (int)strtol(p, &end, 10);
        if (!end)
                return -1;

        /*
         * pid does not match mmap pid
         * pid==0 in system-wide mode (synthesized)
         */
        if (pid && pid2 != pid)
                return -1;
        /*
         * validate suffix
         */
        if (strcmp(end, ".dump"))
                return -1;

        if (verbose > 0)
                fprintf(stderr, "jit marker found: %s\n", mmap_name);

        return 0;
}

int
jit_process(struct perf_session *session,
            struct perf_data_file *output,
            struct machine *machine,
            char *filename,
            pid_t pid,
            u64 *nbytes)
{
        struct perf_evsel *first;
        struct jit_buf_desc jd;
        int ret;

        /*
         * first, detect marker mmap (i.e., the jitdump mmap)
         */
        if (jit_detect(filename, pid))
                return 0;

        memset(&jd, 0, sizeof(jd));

        jd.session = session;
        jd.output  = output;
        jd.machine = machine;

        /*
         * track sample_type to compute id_all layout
         * perf sets the same sample type to all events as of now
         */
        first = perf_evlist__first(session->evlist);
        jd.sample_type = first->attr.sample_type;

        *nbytes = 0;

        ret = jit_inject(&jd, filename);
        if (!ret) {
                *nbytes = jd.bytes_written;
                ret = 1;
        }

        return ret;
}