#define NOB_IMPLEMENTATION
#include "nob.h"

// Replicate your compiler and linker flags.  Keep the literals in macros so
// they can be used in file-scope aggregate initializers.  MSVC C rejects
// initializing a static object with the value of another object, even if that
// object is a const char*.
#define CFLAGS_LITERAL "/nologo /W3 /WX /Od /GS- /Gs999999 /Gz /Gy /EHsc- /GR- /Oi- /Z7 /std:c17 /D_CRT_SECURE_NO_WARNINGS /D_CRT_NONSTDC_NO_WARNINGS /Icode\\public /Icode\\common\\include /Icode\\kernel\\include"
#define LWIP_CFLAGS_LITERAL CFLAGS_LITERAL " /Icode\\drivers\\net\\lwip\\include /Icode\\third_party\\lwip\\src\\include /wd4005 /wd4100 /wd4103 /wd4127 /wd4244 /wd4267 /wd4706"
#define LFLAGS_LITERAL "/nologo /NODEFAULTLIB /DEBUG"

const char *CFLAGS = CFLAGS_LITERAL;
const char *LWIP_CFLAGS = LWIP_CFLAGS_LITERAL;
const char *LFLAGS = LFLAGS_LITERAL;

typedef struct Driver_Target {
    const char *name;
    const char *obj_dir;
    const char *sources;
    const char *cflags;
} Driver_Target;

#define DRIVER_TARGET(name_, obj_dir_, sources_) { (name_), (obj_dir_), (sources_), CFLAGS_LITERAL }
#define DRIVER_TARGET_FLAGS(name_, obj_dir_, sources_, cflags_) { (name_), (obj_dir_), (sources_), (cflags_) }

static const Driver_Target STANDALONE_DRIVERS[] = {
    DRIVER_TARGET("ps2",    "obj\\drivers\\ps2",    "code\\drivers\\input\\ps2\\ps2.c code\\drivers\\input\\ps2\\ps2_i8042.c"),
    DRIVER_TARGET("ide",    "obj\\drivers\\ide",    "code\\drivers\\storage\\ide\\ide.c code\\drivers\\storage\\ide\\ide_bus.c"),
    DRIVER_TARGET("ahci",   "obj\\drivers\\ahci",   "code\\drivers\\storage\\ahci\\ahci.c code\\drivers\\storage\\ahci\\ahci_port.c"),
    DRIVER_TARGET("fat32",  "obj\\drivers\\fat32",  "code\\drivers\\fs\\fat32\\fat32.c code\\drivers\\fs\\fat32\\fat32_mount.c"),
    DRIVER_TARGET("usb",    "obj\\drivers\\usb",    "code\\drivers\\bus\\usb\\usb.c code\\drivers\\bus\\usb\\usb_bus.c code\\drivers\\bus\\usb\\usb_composite.c code\\drivers\\bus\\usb\\usb_core.c code\\drivers\\bus\\usb\\usb_xhci.c code\\drivers\\bus\\usb\\usb_hid.c code\\drivers\\bus\\usb\\usb_storage.c code\\drivers\\bus\\usb\\usb_hub.c"),
    DRIVER_TARGET("rtl816x", "obj\\drivers\\rtl816x", "code\\drivers\\net\\rtl816x\\rtl816x.c code\\drivers\\net\\rtl816x\\rtl816x_pci.c"),
    DRIVER_TARGET("rtl8139", "obj\\drivers\\rtl8139", "code\\drivers\\net\\rtl8139\\rtl8139.c code\\drivers\\net\\rtl8139\\rtl8139_pci.c"),
};

void create_layout() {
    const char *dirs[] = {
        "obj", "bin", "obj\\boot", "obj\\kernel", "obj\\lwip", "obj\\hal",
        "obj\\syscall", "obj\\runtime", "obj\\win32", "obj\\posix", "obj\\drivers", "obj\\init",
        "obj\\shell_win32", "obj\\shell_posix", "obj\\shell_native",
        "obj\\async_apc", "obj\\async_cancel", "obj\\edit", "obj\\httpget",
        "obj\\tcpserve", "obj\\udp", "obj\\resolve", "obj\\ping", "obj\\doom", "obj\\video"
    };
    for (size_t i = 0; i < NOB_ARRAY_LEN(dirs); i++) nob_mkdir_if_not_exists(dirs[i]);
    for (size_t i = 0; i < NOB_ARRAY_LEN(STANDALONE_DRIVERS); i++) nob_mkdir_if_not_exists(STANDALONE_DRIVERS[i].obj_dir);
}

// Helper to run a shell command string and wait for it
bool run_sync(const char *cmd_str) {
    Nob_Cmd cmd = {0};
    nob_cmd_append(&cmd, "cmd.exe", "/c", cmd_str);
    bool ok = nob_cmd_run_sync(cmd);
    nob_cmd_free(cmd);
    return ok;
}

// Helper to launch a shell command into the background
Nob_Proc run_async(const char *cmd_str) {
    Nob_Cmd cmd = {0};
    nob_cmd_append(&cmd, "cmd.exe", "/c", cmd_str);
    Nob_Proc proc = nob_cmd_run_async(cmd);
    nob_cmd_free(cmd);
    return proc;
}

int main(int argc, char **argv) {
    // Magic macro: Checks if build.c is newer than build.exe and recompiles itself if so
    NOB_GO_REBUILD_URSELF(argc, argv);

    // 1. Initial Checks & Setup
    if (!nob_file_exists("disk.vhd")) {
        nob_log(NOB_INFO, "disk.vhd missing, calling mkimage.bat...");
        if (!run_sync("mkimage.bat")) return 1;
    }
    create_layout();

    // 2. Sequential Phase (Bootloader, Core Kernel, and DLLs)
    nob_log(NOB_INFO, "--- Building Core Subsystems ---");
    
    const char *core_tasks[] = {
        "ml64 /nologo /c /Foobj\\kernel\\ke_switch.obj code\\kernel\\ke_switch.asm",
        "ml64 /nologo /c /Foobj\\kernel\\ke_syscall_asm.obj code\\kernel\\ke_syscall.asm",
        "ml64 /nologo /c /Foobj\\kernel\\ke_gdt_reload.obj code\\kernel\\ke_gdt_reload.asm",
        "ml64 /nologo /c /Foobj\\syscall\\stubs.obj code\\user\\syscall\\syscall_stubs.asm",
        "for %f in (code\\boot\\uefi\\*.asm) do ml64 /nologo /c /Fo\"obj\\boot\\%~nf.obj\" \"%f\"",
        
        nob_temp_sprintf("cl %s /MP /c /Foobj\\ code\\common\\printf.c code\\common\\ini_parser.c", CFLAGS),
        "cl /nologo /c /Foobj\\libc.obj code\\common\\libc.c",
        nob_temp_sprintf("cl %s /c /Foobj\\runtime\\user_heap.obj code\\user\\runtime\\heap.c", CFLAGS),
        nob_temp_sprintf("cl %s /c /Foobj\\runtime\\user_heap_exports.obj code\\user\\runtime\\heap_exports.c", CFLAGS),
        
        // Bootloader
        nob_temp_sprintf("cl %s /MP /c /Icode\\boot\\uefi\\include /Foobj\\boot\\ code\\boot\\uefi\\*.c && link %s /FORCE:MULTIPLE /SUBSYSTEM:EFI_APPLICATION /ENTRY:UefiMain /OUT:bin\\bootx64.efi obj\\libc.obj obj\\ini_parser.obj code\\boot\\uefi\\efilib.lib obj\\boot\\*.obj", CFLAGS, LFLAGS),
        
        // HAL
        nob_temp_sprintf("cl %s /MP /c /Foobj\\hal\\ code\\kernel\\hal\\*.c && link %s /DLL /OUT:bin\\hal.dll /ENTRY:HalInitialize /IGNORE:4216 obj\\printf.obj obj\\libc.obj obj\\hal\\*.obj", CFLAGS, LFLAGS),
        
        // Kernel
        nob_temp_sprintf("cl %s /MP /c /Foobj\\kernel\\ code\\kernel\\*.c && link %s /SUBSYSTEM:NATIVE /ENTRY:KernelMain /OUT:bin\\kernel.exe bin\\hal.lib obj\\ini_parser.obj obj\\printf.obj obj\\libc.obj obj\\kernel\\*.obj", CFLAGS, LFLAGS),
        
        // Subsystem DLLs
        "link /nologo /SUBSYSTEM:NATIVE /NODEFAULTLIB /DLL /OUT:bin\\syscall.dll obj\\syscall\\stubs.obj",
        nob_temp_sprintf("cl %s /DUSER_HEAP_EXPORTS /MP /c /Foobj\\win32\\ code\\user\\runtime\\heap.c code\\user\\runtime\\heap_exports.c code\\user\\personalities\\win32\\*.c && link /nologo /SUBSYSTEM:NATIVE /NODEFAULTLIB /DLL /OUT:bin\\win32.dll bin\\syscall.lib obj\\libc.obj obj\\win32\\*.obj", CFLAGS),
        nob_temp_sprintf("cl %s /DUSER_HEAP_EXPORTS /MP /c /Foobj\\posix\\ code\\user\\runtime\\heap.c code\\user\\runtime\\heap_exports.c code\\user\\personalities\\posix\\*.c && link /nologo /SUBSYSTEM:NATIVE /NODEFAULTLIB /DLL /OUT:bin\\posix.dll bin\\syscall.lib obj\\libc.obj obj\\posix\\*.obj", CFLAGS)
    };

    for (size_t i = 0; i < NOB_ARRAY_LEN(core_tasks); i++) {
        if (!run_sync(core_tasks[i])) return 1;
    }

    // 3. Parallel Phase (Drivers and Applications)
    nob_log(NOB_INFO, "--- Compiling Drivers and Apps in Parallel ---");
    Nob_Procs procs = {0};

    // Queue standalone drivers. Each driver is now a target with one or more
    // translation units, but still links to one .sys image with DriverEntry.
    for (size_t i = 0; i < NOB_ARRAY_LEN(STANDALONE_DRIVERS); i++) {
        const Driver_Target *driver = &STANDALONE_DRIVERS[i];
        const char *cmd = nob_temp_sprintf(
            "cl %s /MP /c /Fo%s\\ %s && link /nologo /SUBSYSTEM:NATIVE /NODEFAULTLIB /DLL /OUT:bin\\%s.sys /ENTRY:DriverEntry bin\\hal.lib bin\\kernel.lib obj\\libc.obj %s\\*.obj",
            driver->cflags, driver->obj_dir, driver->sources, driver->name, driver->obj_dir
        );
        nob_da_append(&procs, run_async(cmd));
    }

    // Parse lwIP Sources and Queue it
    Nob_String_Builder sb = {0}, lwip_files = {0};
    if (nob_read_entire_file("code\\drivers\\net\\lwip\\lwip_sources.txt", &sb)) {
        Nob_String_View sv = nob_sv_from_parts(sb.items, sb.count);
        while (sv.count > 0) {
            Nob_String_View line = nob_sv_trim(nob_sv_chop_by_delim(&sv, '\n'));
            if (line.count > 0 && line.data[0] != '#') {
                nob_sb_append_cstr(&lwip_files, "code\\");
                for (size_t i = 0; i < line.count; i++) nob_sb_append_buf(&lwip_files, &(char){line.data[i] == '/' ? '\\' : line.data[i]}, 1);
                nob_sb_append_cstr(&lwip_files, " ");
            }
        }
        nob_sb_append_null(&lwip_files);
        nob_da_append(&procs, run_async(nob_temp_sprintf("cl %s /MP /c /Foobj\\lwip\\ %s && link /nologo /SUBSYSTEM:NATIVE /NODEFAULTLIB /DLL /OUT:bin\\lwip.sys /ENTRY:DriverEntry bin\\hal.lib bin\\kernel.lib obj\\libc.obj obj\\lwip\\*.obj", LWIP_CFLAGS, lwip_files.items)));
    }

    // Queue Subsystem Applications
    const char *apps[][4] = {
        // Folder, Libs, Name, Optional Extra Link Args
        {"system\\init", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "init", "/ENTRY:_start"},
        {"apps\\shell_win32", "bin\\win32.lib obj\\libc.obj", "shell_win32", "/ENTRY:main"},
        {"apps\\shell_posix", "bin\\posix.lib obj\\libc.obj", "shell_posix", "/ENTRY:main"},
        {"apps\\shell_native", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "shell_native", "/ENTRY:main"},
        {"tests\\async_apc", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "async_apc", "/ENTRY:main"},
        {"tests\\async_cancel", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "async_cancel", "/ENTRY:main"},
        {"apps\\edit", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "edit", "/ENTRY:main"},
        {"apps\\httpget", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "httpget", "/ENTRY:main"},
        {"apps\\tcpserve", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "tcpserve", "/ENTRY:main"},
        {"apps\\udp", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "udp", "/ENTRY:main"},
        {"apps\\resolve", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "resolve", "/ENTRY:main"},
        {"apps\\ping", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "ping", "/ENTRY:main"},
        {"apps\\doom", "bin\\syscall.lib obj\\libc.obj obj\\runtime\\user_heap.obj obj\\runtime\\user_heap_exports.obj", "doom", "/ENTRY:main"},
        {"apps\\video", "bin\\win32.lib obj\\libc.obj code\\user\\apps\\video\\bink2w64.lib", "video", "/ENTRY:main"}
    };

    for (size_t i = 0; i < NOB_ARRAY_LEN(apps); i++) {
        const char *cmd = nob_temp_sprintf(
            "cl %s /MP /c /Foobj\\%s\\ code\\user\\%s\\*.c && link %s /SUBSYSTEM:NATIVE %s /OUT:bin\\%s.exe /PDB:bin\\%s.pdb %s obj\\%s\\*.obj",
            CFLAGS, apps[i][2], apps[i][0], LFLAGS, apps[i][3], apps[i][2], apps[i][2], apps[i][1], apps[i][2]
        );
        nob_da_append(&procs, run_async(cmd));
    }

    // Wait for everything to link simultaneously
    if (!nob_procs_wait(procs)) {
        nob_log(NOB_ERROR, "One or more parallel build jobs failed!");
        return 1;
    }

    // 4. Mounting & Packing Image
    nob_log(NOB_INFO, "--- Mounting disk.vhd ---");
    
    // Resolve the absolute path to disk.vhd for diskpart
    char vhd_path[MAX_PATH];
    GetFullPathNameA("disk.vhd", MAX_PATH, vhd_path, NULL);

    const char *diskpart_mnt = nob_temp_sprintf(
        "select vdisk file=\"%s\"\nattach vdisk\n"
        "select partition 1\nformat fs=fat32 quick\nassign letter=Z\n"
        "select partition 2\nformat fs=fat32 quick\nassign letter=Y\nexit\n",
        vhd_path
    );
    nob_write_entire_file("mnt.txt", diskpart_mnt, strlen(diskpart_mnt));
    if (!run_sync("diskpart /s mnt.txt")) return 1;

    nob_log(NOB_INFO, "--- Assembling ESP ---");
    _mkdir("Z:\\apps");     _mkdir("Y:\\apps");
    _mkdir("Z:\\system");   _mkdir("Z:\\drivers");
    _mkdir("Z:\\EFI");      _mkdir("Z:\\EFI\\BOOT");

    const char *copies[][2] = {
        {"bin\\bootx64.efi", "Z:\\EFI\\BOOT\\bootx64.efi"}, {"bin\\kernel.exe", "Z:\\system\\colonel.exe"},
        {"bin\\hal.dll", "Z:\\system\\hal.dll"},            {"bin\\syscall.dll", "Z:\\system\\syscall.dll"},
        {"bin\\win32.dll", "Z:\\system\\win32.dll"},        {"bin\\posix.dll", "Z:\\system\\posix.dll"},
        {"code\\settings.ini", "Z:\\system\\settings.ini"}, {"bin\\init.exe", "Z:\\system\\init.exe"},
        
        {"bin\\ps2.sys", "Z:\\drivers\\ps2.sys"},           {"bin\\ide.sys", "Z:\\drivers\\ide.sys"},
        {"bin\\ahci.sys", "Z:\\drivers\\ahci.sys"},
        {"bin\\fat32.sys", "Z:\\drivers\\fat32.sys"},       {"bin\\lwip.sys", "Z:\\drivers\\lwip.sys"},
        {"bin\\usb.sys", "Z:\\drivers\\usb.sys"},           {"bin\\rtl816x.sys", "Z:\\drivers\\rtl816x.sys"},
        {"bin\\rtl8139.sys", "Z:\\drivers\\rtl8139.sys"},   {"code\\drivers\\drivers.ini", "Z:\\drivers\\drivers.ini"},
        {"code\\boot.ini", "Z:\\boot.ini"},

        {"bin\\shell_win32.exe", "Z:\\apps\\shell_win32.exe"}, {"bin\\shell_posix.exe", "Z:\\apps\\shell_posix.exe"},
        {"bin\\shell_native.exe", "Z:\\apps\\shell_native.exe"}, {"bin\\edit.exe", "Z:\\apps\\edit.exe"},
        {"bin\\httpget.exe", "Z:\\apps\\httpget.exe"},         {"bin\\tcpserve.exe", "Z:\\apps\\tcpserve.exe"},
        {"bin\\udp.exe", "Z:\\apps\\udp.exe"},                 {"bin\\resolve.exe", "Z:\\apps\\resolve.exe"},
        {"bin\\ping.exe", "Z:\\apps\\ping.exe"},               {"bin\\doom.exe", "Z:\\apps\\doom.exe"},
        {"bin\\video.exe", "Z:\\apps\\video.exe"},

        {"bin\\async_apc.exe", "Y:\\apps\\async_apc.exe"},     {"bin\\async_cancel.exe", "Y:\\apps\\async_cancel.exe"}
    };

    for (size_t i = 0; i < NOB_ARRAY_LEN(copies); i++) nob_copy_file(copies[i][0], copies[i][1]);

    if (nob_file_exists("code\\user\\apps\\video\\bink2w64.dll")) nob_copy_file("code\\user\\apps\\video\\bink2w64.dll", "Z:\\apps\\bink2w64.dll");
    if (nob_file_exists("code\\user\\apps\\video\\video.bk2"))    nob_copy_file("code\\user\\apps\\video\\video.bk2", "Z:\\apps\\video.bk2");
    if (nob_file_exists("code\\user\\apps\\doom\\doom1.wad"))     nob_copy_file("code\\user\\apps\\doom\\doom1.wad", "Z:\\apps\\doom1.wad");

    nob_log(NOB_INFO, "--- Unmounting disk.vhd ---");
    const char *diskpart_umnt = nob_temp_sprintf(
        "select vdisk file=\"%s\"\ndetach vdisk\nexit\n",
        vhd_path
    );
    nob_write_entire_file("umnt.txt", diskpart_umnt, strlen(diskpart_umnt));
    run_sync("diskpart /s umnt.txt");

    remove("mnt.txt");
    remove("umnt.txt");

    nob_log(NOB_INFO, "Build complete!");
    return 0;
}