#include <string.h>
#include <stddef.h>
#include "../include/include.h"
#include "include/contract.h"
#include "include/memory.h"

#define MZ           0x5A4Du
#define DIR_EXPORT   0u
#define DIR_DEBUG    6u
#define DBG_CODEVIEW 2u
#define CV_RSDS      0x53445352u

static int beacon_find(gva_ctx* ctx, uint64_t* pa, uint64_t* va) {
    void *ptr = p_(ctx).pa;
    const void *end = p_(ctx).pa + p_(ctx).fsize;

    do {
        const contract* c = (void*)ptr;
        if (c->magic0 == CONTRACT_MAGIC0 && c->magic1 == CONTRACT_MAGIC1) {
            *pa = offset_gpa(&p_(ctx), ptr - p_(ctx).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(gva_ctx* ctx, uintptr_t cr3, uint64_t base, unsigned idx, uint32_t* rva, uint32_t* size) {
    uint32_t lfanew;
    if (gva_read(ctx, cr3, base + 0x3C, &lfanew, 4)) {
        return -1;
    }
    const uint64_t dd = base + lfanew + 0x18 + 0x70 + (uint64_t)idx*8;
    if (gva_read(ctx, cr3, dd, rva, 4)) {
        return -1;
    }
    return (size && gva_read(ctx, cr3, dd + 4, size, 4)) ? -1 : 0;
}

static int pe_pdb(gva_ctx* ctx, 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(ctx, 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(ctx, cr3, base + dbg_rva + o + 0x0C, &type, 4)) {
            return -1;
        }
        if (type != DBG_CODEVIEW) {
            continue;
        }
        if (gva_read(ctx, cr3, base + dbg_rva + o + 0x14, &cv_rva, 4)) { /* AddressOfRawData RVA */
            return -1;
        }
        if (gva_read(ctx, cr3, base + cv_rva, &sig, 4) || sig != CV_RSDS) {
            return -1;
        }
        if (gva_read(ctx, cr3, base + cv_rva + 0x04, guid, 16)) {
            return -1;
        }
        if (gva_read(ctx, cr3, base + cv_rva + 0x14, age, 4)) {
            return -1;
        }
        gva_read(ctx, cr3, base + cv_rva + 0x18, name, namecap); /* best-effort */
        name[namecap - 1] = 0;
        return 0;
    }
    return -1;
}

static int find_ntoskrnl(gva_ctx* ctx, 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(&p_(ctx), 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(&p_(ctx), 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(&p_(ctx), 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(ctx, cr3, va, &mz, 2) || mz != MZ) {
                    continue;
                }
                if (pe_pdb(ctx, 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(gva_ctx* ctx, uintptr_t cr3, uint64_t kbase, const char* want) {
    uint32_t exp_rva;
    if (pe_datadir(ctx, cr3, kbase, DIR_EXPORT, &exp_rva, NULL) || !exp_rva) {
        return 0;
    }

    uint8_t ed[40];
    if (gva_read(ctx, 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(ctx, cr3, kbase + a_names + i*4, &nrva, 4)) {
            return 0;
        }
        if (gva_read(ctx, 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(ctx, cr3, kbase + a_ords + i*2, &ord, 2)) {
            return 0;
        }
        return gva_read(ctx, cr3, kbase + a_funcs + ord*4, &frva, 4) ? 0 : frva;
    }
    return 0;
}

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

__attribute__((cold))
int host_bootstrap(gva_ctx* ctx) {
    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(ctx, &anchor_pa, &va_self)) {
        return -1;
    }

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

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

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

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

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

    beacon_ack(ctx, anchor_pa);
    return 0;
}