#include <stdint.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include "core.h"

#define RAM_H   (1ul<<32)
#define PROT_RW (PROT_READ | PROT_WRITE)

static void clean_ctx(vmie_mem* m) {
    memset(m, 0, sizeof(vmie_mem));
    m->fd = -1;
}

/* Resolve GPA `g` over the segment map: on success returns the seg covering it
 * and the file offset of `g`; NULL if `g` falls outside every seg (a hole). The
 * `g - s->gpa < s->len` test is one branch and folds the lower+upper bound. */
__attribute__((hot))
static const gpa_seg* gpa_seg_of(const vmie_mem* m, uint64_t g, uintptr_t* off) {
    for (int i = 0; i < m->nseg; i++) {
        const gpa_seg* s = &m->seg[i];
        const uint64_t rel = g - s->gpa;
        if (rel < s->len) {
            *off = (uintptr_t)(s->file_off + rel);
            return s;
        }
    }
    return NULL;
}

/* `*offs` (a GPA) resolves to an in-file offset AND the whole [*, *+nmemb) range
 * fits within its single seg. `nmemb > s->len - rel` is at once the file-bounds
 * check and the no-straddle/seam reject (a range may not cross a seg boundary). */
__attribute__((hot))
static int out_of_bounds(vmie_mem* m, uintptr_t* offs, const size_t nmemb) {
    const uint64_t g = *offs;
    const gpa_seg* s = gpa_seg_of(m, g, offs);
    return !s || nmemb > s->len - (g - s->gpa);
}

__attribute__((hot))
int gpa_read(vmie_mem* m, uintptr_t offs, void* buf, const size_t nmemb) {
    if (out_of_bounds(m, &offs, nmemb)) {
        return -1;
    }
    memcpy(buf, m->pa + offs, nmemb);
    return 0;
}

int gpa_write(vmie_mem* m, uintptr_t offs, const void* src, const size_t nmemb) {
    if (m->ro) {
        return -1;
    }
    if (out_of_bounds(m, &offs, nmemb)) {
        return -1;
    }
    memcpy(m->pa + offs, src, nmemb);
    return 0;
}

/* Zero-copy host pointer to [offs, offs+nmemb) GPA, or NULL if that range is not
 * fully backed by the mapped image. Same split + bounds check as gpa_read. */
__attribute__((hot))
void* gpa_ptr(vmie_mem* m, uintptr_t offs, const size_t nmemb) {
    if (out_of_bounds(m, &offs, nmemb)) {
        return NULL;
    }
    return (uint8_t*)m->pa + offs;
}

/* segment table is well-formed against fsize: nonempty, in range, sorted, dense
 * (each seg starts where the previous file span ended), all spans in-file. */
static int segs_valid(const gpa_seg* segs, int nseg, uint64_t fsize) {
    if (nseg < 1 || nseg > VMIE_MAX_SEGS) {
        return 0;
    }
    uint64_t foff = 0;
    for (int i = 0; i < nseg; i++) {
        if (segs[i].file_off != foff
            || segs[i].len == 0
            || segs[i].len > fsize - foff
            || (i > 0 && segs[i].gpa < segs[i - 1].gpa + segs[i - 1].len)) {
            return 0;
        }
        foff += segs[i].len;
    }
    return 1;
}

/* Finish a vmie_mem over an ALREADY-OWNED fd (from open() or dup()): fstat it,
 * validate the seg map against its real size, mmap it shared (PROT_READ when
 * `ro`, else PROT_RW), and fill `m`. On ANY failure close `fd` and clean `m`
 * (no partial state); on success store the fd in `m->fd`, record `ro`, and
 * return 0. The fd's lifetime is now owned by `m` (closed in gpa_close). The
 * single mmap/validate site shared by every constructor. */
__attribute__((cold))
static int gpa_map_owned_fd(vmie_mem* m, int fd, const gpa_seg* segs, int nseg, int ro) {
    struct stat st;

    if (fstat(fd, &st) || !segs_valid(segs, nseg, (uint64_t)st.st_size)) {
        goto close_;
    }

    const int prot = ro ? PROT_READ : PROT_RW;
    if ((m->pa = mmap(NULL, st.st_size, prot, MAP_SHARED, fd, 0)) == MAP_FAILED) {
      close_:
        close(fd);
        clean_ctx(m);
        return -1;
    }

    m->fd    = fd;
    m->fsize = st.st_size;
    m->nseg  = nseg;
    m->ro    = ro;
    memcpy(m->seg, segs, (size_t)nseg * sizeof *segs);

    return 0;
}

/* Build the classic single-`low` QEMU map into out[0..nseg) and return nseg.
 * Below the 4 GiB PCI hole the file maps 1:1 ([0,low)->file[0,low)); at and
 * above 4 GiB it resumes at file offset low. When low >= fsize the hole is never
 * reached, so one inert identity seg covering the whole image suffices (nseg=1).
 * The single low-map site shared by gpa_open (fsize from stat) and gpa_from_fd
 * (fsize from fstat). */
__attribute__((cold))
static int low_segs(uint64_t fsize, uint64_t low, gpa_seg out[2]) {
    if (low >= fsize) {
        out[0] = (gpa_seg){ 0, fsize, 0 };
        return 1;
    }
    out[0] = (gpa_seg){ 0, low, 0 };
    out[1] = (gpa_seg){ RAM_H, fsize - low, low };
    return 2;
}

__attribute__((cold))
int gpa_open_segs(vmie_mem* m, const char* path, const gpa_seg* segs, int nseg) {
    const int fd = open(path, O_RDWR);
    if (fd < 0) {
        clean_ctx(m);
        return -1;
    }
    return gpa_map_owned_fd(m, fd, segs, nseg, 0);
}

__attribute__((cold))
int gpa_open(vmie_mem* m, const char* path, uintptr_t low) {
    struct stat st;
    if (stat(path, &st)) {
        clean_ctx(m);
        return -1;
    }

    gpa_seg   segs[2];
    const int nseg = low_segs((uint64_t)st.st_size, low, segs);

    return gpa_open_segs(m, path, segs, nseg);
}

/* fd-backed constructors: build a vmie_mem from an already-open file descriptor
 * instead of a path. The fd is BORROWED via dup() - core owns the copy (closed
 * in gpa_close); the caller's fd stays valid. The fd must reference an mmap-able
 * RW object (regular file / memfd / shm); its file position is not used. */
__attribute__((cold))
int gpa_from_fd_segs(vmie_mem* m, int fd, const gpa_seg* segs, int nseg) {
    const int dupfd = dup(fd);
    if (dupfd < 0) {
        clean_ctx(m);
        return -1;
    }
    return gpa_map_owned_fd(m, dupfd, segs, nseg, 0);
}

__attribute__((cold))
int gpa_from_fd(vmie_mem* m, int fd, uintptr_t low) {
    struct stat st;
    if (fstat(fd, &st)) {
        clean_ctx(m);
        return -1;
    }

    gpa_seg   segs[2];
    const int nseg = low_segs((uint64_t)st.st_size, low, segs);

    return gpa_from_fd_segs(m, fd, segs, nseg);
}

/* read-only twins of gpa_from_fd*: map the borrowed fd PROT_READ and mark the
 * source read-only (m->ro), so gpa_write/gva_write return -1. Same dup-borrow
 * semantics (core owns the copy, caller's fd stays valid). Accept O_RDONLY as
 * well as O_RDWR fds - only read access is required for a PROT_READ mapping.
 * Reads/gpa_ptr/scan behave exactly as on an RW source. */
__attribute__((cold))
int gpa_from_ro_fd_segs(vmie_mem* m, int fd, const gpa_seg* segs, int nseg) {
    const int dupfd = dup(fd);
    if (dupfd < 0) {
        clean_ctx(m);
        return -1;
    }
    return gpa_map_owned_fd(m, dupfd, segs, nseg, 1);
}

__attribute__((cold))
int gpa_from_ro_fd(vmie_mem* m, int fd, uintptr_t low) {
    struct stat st;
    if (fstat(fd, &st)) {
        clean_ctx(m);
        return -1;
    }

    gpa_seg   segs[2];
    const int nseg = low_segs((uint64_t)st.st_size, low, segs);

    return gpa_from_ro_fd_segs(m, fd, segs, nseg);
}

__attribute__((cold))
void gpa_close(vmie_mem* m) {
    if (m->pa) {
        munmap(m->pa, m->fsize);
    }

    if (m->fd >= 0) {
        close(m->fd);
    }

    clean_ctx(m);
}

/* ---- public dump source (heap-owned vmie_mem) ---------------------------- *
 * Thin wrappers over gpa_open*: heap-allocate a vmie_mem and open into it, so a
 * dump (or any flat/segmented RAM image) is a first-class memory source for the
 * physical scanners without exposing the win32 engine. */

__attribute__((cold))
vmie_mem* vmie_mem_open(const char* path, uint64_t low) {
    vmie_mem* m = calloc(1, sizeof *m);
    if (!m) {
        return NULL;
    }
    if (gpa_open(m, path, (uintptr_t)low)) {
        free(m);
        return NULL;
    }
    return m;
}

__attribute__((cold))
vmie_mem* vmie_mem_open_segs(const char* path, const gpa_seg* segs, int nseg) {
    vmie_mem* m = calloc(1, sizeof *m);
    if (!m) {
        return NULL;
    }
    if (gpa_open_segs(m, path, segs, nseg)) {
        free(m);
        return NULL;
    }
    return m;
}

__attribute__((cold))
vmie_mem* vmie_mem_from_fd(int fd, uint64_t low) {
    vmie_mem* m = calloc(1, sizeof *m);
    if (!m) {
        return NULL;
    }
    if (gpa_from_fd(m, fd, (uintptr_t)low)) {
        free(m);
        return NULL;
    }
    return m;
}

__attribute__((cold))
vmie_mem* vmie_mem_from_fd_segs(int fd, const gpa_seg* segs, int nseg) {
    vmie_mem* m = calloc(1, sizeof *m);
    if (!m) {
        return NULL;
    }
    if (gpa_from_fd_segs(m, fd, segs, nseg)) {
        free(m);
        return NULL;
    }
    return m;
}

__attribute__((cold))
vmie_mem* vmie_mem_from_ro_fd(int fd, uint64_t low) {
    vmie_mem* m = calloc(1, sizeof *m);
    if (!m) {
        return NULL;
    }
    if (gpa_from_ro_fd(m, fd, (uintptr_t)low)) {
        free(m);
        return NULL;
    }
    return m;
}

__attribute__((cold))
vmie_mem* vmie_mem_from_ro_fd_segs(int fd, const gpa_seg* segs, int nseg) {
    vmie_mem* m = calloc(1, sizeof *m);
    if (!m) {
        return NULL;
    }
    if (gpa_from_ro_fd_segs(m, fd, segs, nseg)) {
        free(m);
        return NULL;
    }
    return m;
}

__attribute__((cold))
void vmie_mem_close(vmie_mem* m) {
    if (!m) {
        return;
    }
    gpa_close(m);
    free(m);
}