#include <stdint.h>
#include <string.h>
#include "../include/include.h"
#include "include/memory.h"

#define pr_(ctx) ((ctx)->prof)

#define RING_CAP 4096   /* USER_MIN/USER_MAX/KERN_MIN come from include/memory.h */
#define SCAN_MAX 1024
#define FT_LO    0x01D0000000000000ll
#define FT_HI    0x01F0000000000000ll

static int canon_ok(uint64_t p, int kernel) {
    return kernel ? (p >= KERN_MIN) : (p >= USER_MIN && p <= USER_MAX);
}

/* Circular LIST_ENTRY walker (Flink at node+0); one primitive for both rings. */
static int list_ring_ok(gva_ctx* ctx, uintptr_t cr3, uint64_t head, int kernel) {
    uint64_t node;
    if (gva_read(ctx, cr3, head, &node, 8)) {
        return 0;
    }
    for (int i = 0; i < RING_CAP; i++) {
        if (node == head) {
            return i > 0;
        }
        if (!canon_ok(node, kernel) || gva_read(ctx, cr3, node, &node, 8)) {
            return 0;
        }
    }
    return 0;
}

/* Pass 1: ep_name/ep_pid/ep_links/ep_dtb from the System _EPROCESS. */
static int discover_core(gva_ctx* ctx, uintptr_t cr3, uint64_t sys_ep) {
    profile* p = &pr_(ctx);
    uint8_t buf[0x800];
    if (gva_read(ctx, cr3, sys_ep, buf, sizeof buf)) {
        return -1;
    }

    int name_off = -1;
    for (int o = 0x100; o + 7 <= (int)sizeof buf; o++) {
        if (!memcmp(buf + o, "System", 6) && buf[o + 6] == 0) {
            name_off = o;
            break;
        }
    }
    if (name_off < 0) {
        return -2;
    }
    p->ep_name = (uint16_t)name_off;

    int pid_off = -1;
    for (int o = 0x80; o + 8 <= name_off; o += 8) {
        if (*(uint64_t*)(buf + o) != 4) {
            continue;
        }
        const uint16_t links = (uint16_t)(o + 8);
        if (list_ring_ok(ctx, cr3, sys_ep + links, 1)) {
            p->ep_links = links;
            pid_off = o;
            break;
        }
    }
    if (pid_off < 0) {
        return -3;
    }
    p->ep_pid = (uint16_t)pid_off;

    int dtb_off = -1;
    for (int o = 0x18; o <= 0x60; o += 8) {
        const uint64_t c = *(uint64_t*)(buf + o) & PFN_MASK;
        uint8_t probe;
        if (c && khalf_score(ctx, c) >= 16 && !gva_read(ctx, c, sys_ep, &probe, 1)) {
            dtb_off = o;
            break;
        }
    }
    if (dtb_off < 0) {
        return -4;
    }
    p->ep_dtb = (uint16_t)dtb_off;
    return 0;
}

/* Transient snapshot of (eprocess, pid, cr3) over the active ring. */
static int collect_procs(gva_ctx* ctx, uintptr_t cr3, uint64_t sys_ep, uint64_t* eps, uint32_t* pids, uint64_t* cr3s, int cap) {
    const profile* p = &pr_(ctx);
    int n = 0;
    uint64_t ep = sys_ep, node;
    do {
        uint64_t pid = 0, dtb = 0;
        gva_read(ctx, cr3, ep + p->ep_pid, &pid, 8);
        gva_read(ctx, cr3, ep + p->ep_dtb, &dtb, 8);
        eps[n]  = ep;
        pids[n] = (uint32_t)pid;
        cr3s[n] = dtb & PFN_MASK;
        if (++n >= cap) {
            break;
        }
        if (gva_read(ctx, cr3, ep + p->ep_links, &node, 8)) {
            break;
        }
        ep = node - p->ep_links;
    } while (ep != sys_ep);
    return n;
}

/* Pass 2a: ep_ppid by population (creator PID). Best-effort. */
static void discover_ppid(gva_ctx* ctx, uintptr_t cr3, const uint64_t* eps, const uint32_t* pids, int n) {
    int best_off = -1, best_hits = 0;
    for (int o = 0x100; o <= 0x600; o += 8) {
        int hits = 0;
        for (int i = 0; i < n; i++) {
            uint32_t cand = 0;
            if (gva_read(ctx, cr3, eps[i] + o, &cand, 4) || !cand || cand == pids[i]) {
                continue;
            }
            for (int j = 0; j < n; j++) {
                if (pids[j] == cand) { hits++; break; }
            }
        }
        if (hits > best_hits) {
            best_hits = hits;
            best_off = o;
        }
    }
    if (best_off >= 0 && best_hits * 3 >= n) {
        pr_(ctx).ep_ppid = (uint16_t)best_off;
    }
}

/* Pass 2b: ep_createtime (CreateTime, FILETIME) -- every sample in boot range, System earliest. Best-effort. */
static void discover_createtime(gva_ctx* ctx, uintptr_t cr3, const uint64_t* eps, int n) {
    for (int o = 0x140; o <= 0x600; o += 8) {
        int64_t sysv = 0;
        int ok = 1;
        for (int i = 0; i < n; i++) {
            int64_t t = 0;
            if (gva_read(ctx, cr3, eps[i] + o, &t, 8) || t < FT_LO || t > FT_HI) { ok = 0; break; }
            if (i == 0) {
                sysv = t;
            } else if (t < sysv) {
                ok = 0; break;
            }
        }
        if (ok) {
            pr_(ctx).ep_createtime = (uint16_t)o;
            return;
        }
    }
}

/* Pass 2c: ep_imgpath (ImageFilePathHint) -- UNICODE_STRING whose tail equals the
 * process's untruncated ImageFileName; probe short-named (<15) procs only. Best-effort. */
static void discover_imgpath(gva_ctx* ctx, uintptr_t cr3, const uint64_t* eps, const uint64_t* cr3s, int n) {
    profile* p = &pr_(ctx);
    for (int i = 0; i < n; i++) {
        if (!cr3s[i]) {
            continue;
        }
        char nm[16] = {0};
        if (gva_read(ctx, cr3, eps[i] + p->ep_name, nm, 15)) {
            continue;
        }
        const size_t nl = strnlen(nm, 15);
        if (nl == 0 || nl >= 15) {
            continue;
        }
        for (int o = 0x400; o <= 0x600; o += 8) {
            uint16_t len = 0;
            uint64_t buf = 0;
            if (gva_read(ctx, cr3, eps[i] + o, &len, 2) || gva_read(ctx, cr3, eps[i] + o + 8, &buf, 8)) {
                continue;
            }
            if ((len & 1) || len < (uint16_t)(nl * 2) || len > 0x800 || buf < KERN_MIN) {
                continue;
            }
            uint16_t w[16];
            if (gva_read(ctx, cr3, buf + len - (uint64_t)nl * 2, w, nl * 2)) {
                continue;
            }
            int match = 1;
            for (size_t c = 0; c < nl; c++) {
                if ((w[c] < 0x80 ? (char)w[c] : 0) != nm[c]) { match = 0; break; }
            }
            if (match) {
                p->ep_imgpath = (uint16_t)o;
                return;
            }
        }
    }
}

/* Pass 2d: ep_peb + user PEB/Ldr chain; commits the x64-invariant LDR offsets
 * (incl. FullDllName) after validating them on the live first entry. */
static int discover_user_chain(gva_ctx* ctx, uintptr_t cr3, const uint64_t* eps, const uint64_t* cr3s, int n) {
    profile* p = &pr_(ctx);

    for (int i = 0; i < n; i++) {
        const uint64_t pcr3 = cr3s[i];
        if (!pcr3) {
            continue;
        }
        for (int po = 0x280; po <= 0x580; po += 8) {
            uint64_t peb = 0;
            if (gva_read(ctx, cr3, eps[i] + po, &peb, 8) || !canon_ok(peb, 0)) {
                continue;
            }
            for (int lo = 0x10; lo <= 0x30; lo += 8) {
                uint64_t ldr = 0;
                if (gva_read(ctx, pcr3, peb + lo, &ldr, 8) || !canon_ok(ldr, 0)) {
                    continue;
                }
                for (int ll = 0x10; ll <= 0x20; ll += 8) {
                    if (!list_ring_ok(ctx, pcr3, ldr + ll, 0)) {
                        continue;
                    }

                    uint64_t entry = 0, dllbase = 0, bufp = 0, fbufp = 0;
                    uint16_t nlen = 0, flen = 0;
                    if (gva_read(ctx, pcr3, ldr + ll, &entry, 8)) {
                        continue;
                    }
                    if (gva_read(ctx, pcr3, entry + 0x30, &dllbase, 8) ||
                        !canon_ok(dllbase, 0) || (dllbase & 0xFFF)) {
                        continue;
                    }
                    if (gva_read(ctx, pcr3, entry + 0x58, &nlen, 2) || !nlen || (nlen & 1) ||
                        gva_read(ctx, pcr3, entry + 0x58 + 8, &bufp, 8) || !canon_ok(bufp, 0)) {
                        continue;
                    }
                    if (gva_read(ctx, pcr3, entry + 0x48, &flen, 2) || (flen & 1) ||
                        gva_read(ctx, pcr3, entry + 0x48 + 8, &fbufp, 8) || !canon_ok(fbufp, 0)) {
                        continue;
                    }

                    p->ep_peb       = (uint16_t)po;
                    p->peb_ldr      = (uint16_t)lo;
                    p->ldr_loadlist = (uint16_t)ll;
                    p->lde_base     = 0x30;
                    p->lde_size     = 0x40;
                    p->lde_fullname = 0x48;
                    p->lde_name     = 0x58;
                    return 0;
                }
            }
        }
    }
    return -1;
}

__attribute__((cold))
int profile_build(gva_ctx* ctx, uintptr_t cr3, uint64_t sys_ep, const uint8_t guid[16], uint32_t age) {
    memset(&pr_(ctx), 0, sizeof(pr_(ctx)));
    memcpy(pr_(ctx).guid, guid, 16);
    pr_(ctx).age = age;

    if (discover_core(ctx, cr3, sys_ep)) {
        return -1;
    }

    uint64_t eps[SCAN_MAX], cr3s[SCAN_MAX];
    uint32_t pids[SCAN_MAX];
    const int n = collect_procs(ctx, cr3, sys_ep, eps, pids, cr3s, SCAN_MAX);
    if (n <= 1) {
        return -2;
    }

    discover_ppid(ctx, cr3, eps, pids, n);
    discover_createtime(ctx, cr3, eps, n);
    discover_imgpath(ctx, cr3, eps, cr3s, n);
    if (discover_user_chain(ctx, cr3, eps, cr3s, n)) {
        return -3;
    }
    return 0;
}