#include <string.h>
#include <stddef.h>
#include <stdlib.h>
#include "win32.h"
#include "engine-win32.h"

#define MZ           0x5A4Du

/* ---- lifecycle (cold) ---------------------------------------------------- */

__attribute__((cold))
vmie_win32* vmie_win32_open(const char* ram_path, uint64_t low) {
    vmie_win32* v = calloc(1, sizeof *v);
    if (!v) {
        return NULL;
    }
    if (gpa_open(&v->mem, ram_path, low)) {
        free(v);
        return NULL;
    }
    return v;
}

__attribute__((cold))
void vmie_win32_close(vmie_win32* v) {
    if (!v) {
        return;
    }
    gpa_close(&v->mem);
    free(v);
}

/* ---- bootstrap CR3 recovery (cold) --------------------------------------- *
 * Scan the image for a candidate PML4 that translates the agent's self-VA to its
 * known physical anchor, scored by kernel-half liveness. The page-table walk is
 * borrowed from the arch layer via gva_translate (cold), so the hot path stays
 * private to gva.c. */
__attribute__((cold))
int cr3_recover(vmie_win32* v, uint64_t va_self, uint64_t target_pa, uintptr_t* cr3_out) {
    vmie_mem* m = &v->mem;
    int best_score = -1; uint64_t best = 0;
    for (size_t off = 0; off + 0x1000 <= m->fsize; off += 0x1000) {
        const uintptr_t cand = offset_gpa(m, off);
        uintptr_t gpa;
        if (gva_translate(m, cand, va_self, &gpa)) continue;
        if ((gpa & ~0xFFFull) != (target_pa & ~0xFFFull)) continue;
        const int score = khalf_score(m, cand);
        if (score > best_score) { best_score = score; best = cand; }
    }
    if (best_score < 0) return -1;
    *cr3_out = best;
    return 0;
}
#define DIR_EXPORT   0u
#define DIR_DEBUG    6u
#define DBG_CODEVIEW 2u
#define CV_RSDS      0x53445352u

static int beacon_find(vmie_mem* m, uint64_t* pa, uint64_t* va) {
    void *ptr = m->pa;
    const void *end = m->pa + m->fsize;

    do {
        const contract* c = (void*)ptr;
        if (c->magic0 == CONTRACT_MAGIC0 && c->magic1 == CONTRACT_MAGIC1) {
            *pa = offset_gpa(m, ptr - m->pa);
            *va = c->va_self;
            return 0;
        }

        ptr += 1ull<<12; /* 4KB step: a locked, page-granular beacon */
    } while (ptr < end);

    return -1;
}

static int pe_datadir(vmie_mem* m, uintptr_t cr3, uint64_t base, unsigned idx, uint32_t* rva, uint32_t* size) {
    uint32_t lfanew;
    if (gva_read(m, cr3, base + 0x3C, &lfanew, 4)) {
        return -1;
    }
    const uint64_t dd = base + lfanew + 0x18 + 0x70 + (uint64_t)idx*8;
    if (gva_read(m, cr3, dd, rva, 4)) {
        return -1;
    }
    return (size && gva_read(m, cr3, dd + 4, size, 4)) ? -1 : 0;
}

static int pe_pdb(vmie_mem* m, uintptr_t cr3, uint64_t base, uint8_t guid[16], uint32_t* age, char* name, size_t namecap) {
    uint32_t dbg_rva, dbg_sz;
    if (pe_datadir(m, cr3, base, DIR_DEBUG, &dbg_rva, &dbg_sz) || !dbg_rva) {
        return -1;
    }

    for (uint32_t o = 0; o + 0x1C <= dbg_sz; o += 0x1C) { /* IMAGE_DEBUG_DIRECTORY[] (28B) */
        uint32_t type, cv_rva, sig;
        if (gva_read(m, cr3, base + dbg_rva + o + 0x0C, &type, 4)) {
            return -1;
        }
        if (type != DBG_CODEVIEW) {
            continue;
        }
        if (gva_read(m, cr3, base + dbg_rva + o + 0x14, &cv_rva, 4)) { /* AddressOfRawData RVA */
            return -1;
        }
        if (gva_read(m, cr3, base + cv_rva, &sig, 4) || sig != CV_RSDS) {
            return -1;
        }
        if (gva_read(m, cr3, base + cv_rva + 0x04, guid, 16)) {
            return -1;
        }
        if (gva_read(m, cr3, base + cv_rva + 0x14, age, 4)) {
            return -1;
        }
        gva_read(m, cr3, base + cv_rva + 0x18, name, namecap); /* best-effort */
        name[namecap - 1] = 0;
        return 0;
    }
    return -1;
}

static int find_ntoskrnl(vmie_mem* m, uintptr_t cr3, uint64_t* base, uint8_t guid[16], uint32_t* age) {
    const uint64_t t = cr3 & PFN_MASK;

    for (int p4 = 256; p4 < 512; p4++) {
        uint64_t e4;
        if (gpa_read(m, t + p4*8, &e4, 8) || !(e4 & PG_P)) {
            continue;
        }
        const uint64_t pdpt = e4 & PFN_MASK;

        for (int p3 = 0; p3 < 512; p3++) {
            uint64_t e3;
            if (gpa_read(m, pdpt + p3*8, &e3, 8) || !(e3 & PG_P)) {
                continue;
            }
            if (e3 & PG_PS) {
                continue; /* 1G leaf -- no PE image here */
            }
            const uint64_t pd = e3 & PFN_MASK;

            for (int p2 = 0; p2 < 512; p2++) {
                uint64_t e2;
                if (gpa_read(m, pd + p2*8, &e2, 8) || !(e2 & PG_P)) {
                    continue;
                }

                uint64_t va = (uint64_t)p4<<39 | (uint64_t)p3<<30 | (uint64_t)p2<<21;
                va = VA_CANON(va);

                uint16_t mz; char pdb[16] = {0};
                if (gva_read(m, cr3, va, &mz, 2) || mz != MZ) {
                    continue;
                }
                if (pe_pdb(m, cr3, va, guid, age, pdb, sizeof pdb)) {
                    continue;
                }
                if (strncmp(pdb, "ntkrnlmp.pdb", 12) != 0) {
                    continue;
                }
                *base = va;
                return 0;
            }
        }
    }
    return -1;
}

static uint32_t ko_export_rva(vmie_mem* m, uintptr_t cr3, uint64_t kbase, const char* want) {
    uint32_t exp_rva;
    if (pe_datadir(m, cr3, kbase, DIR_EXPORT, &exp_rva, NULL) || !exp_rva) {
        return 0;
    }

    uint8_t ed[40];
    if (gva_read(m, cr3, kbase + exp_rva, ed, sizeof ed)) {
        return 0;
    }
    uint32_t nnames, a_funcs, a_names, a_ords;
    memcpy(&nnames,  ed + 0x18, 4);
    memcpy(&a_funcs, ed + 0x1C, 4);
    memcpy(&a_names, ed + 0x20, 4);
    memcpy(&a_ords,  ed + 0x24, 4);

    for (uint32_t i = 0; i < nnames; i++) {
        uint32_t nrva; char nm[40];
        if (gva_read(m, cr3, kbase + a_names + i*4, &nrva, 4)) {
            return 0;
        }
        if (gva_read(m, cr3, kbase + nrva, nm, sizeof nm)) {
            continue;
        }
        nm[sizeof nm - 1] = 0;
        if (strcmp(nm, want) != 0) {
            continue;
        }
        uint16_t ord; uint32_t frva;
        if (gva_read(m, cr3, kbase + a_ords + i*2, &ord, 2)) {
            return 0;
        }
        return gva_read(m, cr3, kbase + a_funcs + ord*4, &frva, 4) ? 0 : frva;
    }
    return 0;
}

static void beacon_ack(vmie_mem* m, uint64_t anchor_pa) {
    uint64_t ack = CONTRACT_ACK;
    gpa_write(m, anchor_pa + offsetof(contract, ack), &ack, 8);
}

vmie_mem* vmie_win32_mem(vmie_win32* v) {
    return v ? &v->mem : NULL;
}

__attribute__((cold))
int host_bootstrap(vmie_win32* v) {
    vmie_mem* m = &v->mem;
    uint64_t anchor_pa, va_self;
    uintptr_t cr3boot;
    uint32_t rva;
    uint8_t guid[16];
    uint32_t age;
    uint64_t sys_ep;

    if (beacon_find(m, &anchor_pa, &va_self)) {
        return -1;
    }

    if (cr3_recover(v, va_self, anchor_pa, &cr3boot)) {
        return -2;
    }

    if (find_ntoskrnl(m, cr3boot, &v->kbase, guid, &age)) {
        return -3;
    }

    rva = ko_export_rva(m, cr3boot, v->kbase, "PsInitialSystemProcess");
    if (!rva || gva_read(m, cr3boot, v->kbase + rva, &sys_ep, 8)) {
        return -4;
    }

    if (profile_build(v, cr3boot, sys_ep, guid, age)) {
        return -5;
    }

    uint64_t dtb;
    if (gva_read(m, cr3boot, sys_ep + v->prof.ep_dtb, &dtb, 8)) {
        return -6;
    }
    v->kcr3    = dtb & PFN_MASK;
    v->sysproc = sys_ep;

    beacon_ack(m, anchor_pa);
    return 0;
}