//
// Copyright 2026 Aarav Ravindra Kharade
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

/*
 * ArkOS Build System v3.0
 * -----------------------
 * Orchestrates the complete OS build pipeline from source to bootable images.
 *
 * Architecture:
 *   1. Clean previous output (preserving build caches)
 *   2. Stage frameworks, boot assets, vendor blobs, and service configs
 *   3. Compile all system binaries (init, arkrt, splash, display daemon, UI
 * app)
 *   4. Build the SwiftUI framework via SPM (cached — only rebuilds on changes)
 *   5. Sign the runtime binary for Verified Boot
 *   6. Pack initramfs with all runtime binaries
 *   7. Assemble bootloader and generate disk images
 *
 * All intermediate artifacts go into out_staging/. Final outputs go into
 * finished/. The .swift_build/ directory inside out_staging/ is preserved
 * across builds to enable SPM incremental compilation caching.
 *
 * Usage: build <arkos-base-dir>
 */

#include "ark_parser.h"
#include <dirent.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

/* ── Build Configuration ───────────────────────────────────────── */

#define TOTAL_STEPS 16
#define RPI4_TOTAL_STEPS 10

static int current_step = 0;
static int total_steps = TOTAL_STEPS;

/* ── Progress Reporting ────────────────────────────────────────── */

static void step(const char *desc) {
  current_step++;
  printf("\n\033[1;36m[%2d/%d]\033[0m %s\n", current_step, total_steps, desc);
}

/* ── Command Execution ─────────────────────────────────────────── */

/**
 * Execute a shell command and print it. Non-zero exit codes are reported
 * but do NOT abort the build — callers must check return values for
 * critical steps.
 */
static int run(const char *cmd) {
  printf("  \033[0;90m$ %s\033[0m\n", cmd);
  int ret = system(cmd);
  if (ret != 0) {
    fprintf(stderr, "  \033[1;31m✗ Command failed (exit %d)\033[0m\n", ret);
  }
  return ret;
}

/* ── Path Utilities ────────────────────────────────────────────── */

/** Join two path components. Caller must free the returned string. */
static char *pjoin(const char *a, const char *b) {
  size_t len = strlen(a) + 1 + strlen(b) + 1;
  char *out = malloc(len);
  snprintf(out, len, "%s/%s", a, b);
  return out;
}

/* ── Incremental Build Helpers ─────────────────────────────────── */

/**
 * Check if an output file exists and is newer than all listed source files.
 * Returns 1 if the output is up-to-date (build can be skipped), 0 otherwise.
 *
 * This enables incremental builds — Swift compilation is expensive, so we
 * only recompile when source files have actually changed.
 */
static int is_up_to_date(const char *output, const char **sources,
                         int source_count) {
  struct stat out_st;
  if (stat(output, &out_st) != 0) {
    return 0; /* Output doesn't exist — must build */
  }

  for (int i = 0; i < source_count; i++) {
    struct stat src_st;
    if (stat(sources[i], &src_st) != 0) {
      return 0; /* Source missing — must build */
    }
    if (src_st.st_mtime > out_st.st_mtime) {
      return 0; /* Source is newer — must rebuild */
    }
  }
  return 1; /* All sources are older than output — skip */
}

/* ── RPi4 Build Pipeline ───────────────────────────────────────── */

/* ── Main Build Pipeline ───────────────────────────────────────── */

int main(int argc, char *argv[]) {
  if (argc < 2) {
    fprintf(stderr, "Usage: %s <arkos-base-dir> [--device rpi4] [--no-image]\n",
            argv[0]);
    return 1;
  }

  const char *base = argv[1]; /* e.g. /home/user/repo/arkos */

  int is_rpi4 = 0;
  int no_image =
      0; /* When set, skip merging individual .img files into one rpi4.img */
  for (int i = 2; i < argc; i++) {
    if (strcmp(argv[i], "--device") == 0 && i + 1 < argc) {
      if (strcmp(argv[i + 1], "rpi4") == 0) {
        is_rpi4 = 1;
      } else {
        fprintf(stderr, "Unknown device: %s\n", argv[i + 1]);
        return 1;
      }
    } else if (strcmp(argv[i], "--no-image") == 0) {
      no_image = 1;
    }
  }

  /* Derive the repo root (one level up from the arkos directory) */
  char *repo = pjoin(base, "..");

  /* Output and staging directories */
  char *out_dir =
      is_rpi4 ? pjoin(base, "finished/rpi4") : pjoin(base, "finished");
  char *staging =
      is_rpi4 ? pjoin(base, "out_staging/rpi4") : pjoin(base, "out_staging");
  char *staging_sys = pjoin(staging, "system");
  char *staging_fw = pjoin(staging, "system/frameworks");
  char *staging_vend = pjoin(staging, "vendor");
  char *staging_boot = pjoin(staging, "boot");

  /* Source directories */
  char *fw_dir = pjoin(base, "frameworks");
  char *boot_dir = pjoin(base, "boot");
  char *vendor_dir = pjoin(base, "vendor");
  char *system_dir = pjoin(base, "system");
  char *kernel_dir = pjoin(base, "kernel");

  printf("\n\033[1;35m╔══════════════════════════════════════╗\033[0m\n");
  printf("\033[1;35m║      ArkOS Build System v3.0         ║\033[0m\n");
  if (is_rpi4) {
    printf("\033[1;35m║      Target: ARM64 (Raspberry Pi)    ║\033[0m\n");
  } else {
    printf("\033[1;35m║      Target: x86_64 (PC / QEMU)      ║\033[0m\n");
  }
  printf("\033[1;35m╚══════════════════════════════════════╝\033[0m\n");

  /* ─── Step 1: Clean previous output ────────────────────────── */
  step("Cleaning previous output (preserving build caches)");
  {
    char cmd[1024];
    /* Preserve out_staging/ entirely for incremental builds.
     * Only wipe finished/ which contains the final disk images. */
    snprintf(cmd, sizeof(cmd), "rm -rf %s", out_dir);
    run(cmd);
  }

  /* ─── Step 2: Create staging directories ───────────────────── */
  step("Creating staging directories");
  {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd), "mkdir -p %s %s %s %s %s/services", out_dir,
             staging_fw, staging_vend, staging_boot, staging_sys);
    run(cmd);
  }

  /* ─── Step 3: Stage frameworks ─────────────────────────────── */
  step("Staging frameworks");
  {
    char rm_cmd[1024];
    snprintf(rm_cmd, sizeof(rm_cmd), "rm -rf %s", staging_fw);
    run(rm_cmd);
    char mk_cmd[1024];
    snprintf(mk_cmd, sizeof(mk_cmd), "mkdir -p %s", staging_fw);
    run(mk_cmd);
    DIR *d = opendir(fw_dir);
    if (d) {
      struct dirent *ent;
      while ((ent = readdir(d)) != NULL) {
        if (ent->d_name[0] == '.')
          continue;
        /* DRM framework is excluded from the base system image */
        if (strcmp(ent->d_name, "DRM") == 0) {
          printf("  Skipping excluded framework: %s\n", ent->d_name);
          continue;
        }
        char cmd[1024];
        snprintf(cmd, sizeof(cmd), "cp -ur %s/%s %s/", fw_dir, ent->d_name,
                 staging_fw);
        run(cmd);
      }
      closedir(d);
    }
  }

  /* ─── Step 4: Stage boot assets ────────────────────────────── */
  step("Staging boot assets");
  {
    DIR *d = opendir(boot_dir);
    if (d) {
      struct dirent *ent;
      while ((ent = readdir(d)) != NULL) {
        if (ent->d_name[0] == '.')
          continue;
        char cmd[1024];
        snprintf(cmd, sizeof(cmd), "cp -ur %s/%s %s/", boot_dir, ent->d_name,
                 staging_boot);
        run(cmd);
      }
      closedir(d);
    }
  }

  /* ─── Step 5: Stage system services ────────────────────────── */
  step("Staging system service configs");
  {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd),
             "cp -ur %s/system/services/* %s/services/ 2>/dev/null || true",
             base, staging_sys);
    run(cmd);
  }

  /* ─── Step 6: Stage vendor files ───────────────────────────── */
  step("Staging vendor files");
  {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd), "cp -ur %s/* %s/ 2>/dev/null || true",
             vendor_dir, staging_vend);
    run(cmd);
    /* Exclude private Verified Boot key from public vendor partition */
    snprintf(cmd, sizeof(cmd),
             "rm -f %s/verify/securebuild.ark %s/securebuild.ark 2>/dev/null "
             "|| true",
             staging_vend, staging_vend);
    run(cmd);
  }

  /* ─── Step 7: Verify kernel configuration ──────────────────── */
  step("Verifying kernel configuration");
  {
    char *kconfig = pjoin(kernel_dir, ".config");
    if (access(kconfig, F_OK) != -1) {
      printf("  \033[0;32m✓\033[0m Kernel .config present\n");
    } else {
      printf("  \033[1;33m⚠ WARNING:\033[0m Kernel .config not found!\n");
    }
    free(kconfig);
  }

  /* ─── Step 8: Compile fb_helper.c ──────────────────────────── */
  step(is_rpi4 ? "Cross-compiling fb_helper.c (ARM64)"
               : "Compiling fb_helper.c (framebuffer ioctl bridge)");
  {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd), "make -C %s %s/fb_helper.o", base,
             is_rpi4 ? "out_staging/rpi4" : "out_staging");
    run(cmd);
  }

  /* ─── Step 9: Compile Swift system binaries (incremental) ──── */
  step(is_rpi4
           ? "Cross-compiling Swift system binaries (ARM64)"
           : "Compiling Swift system binaries (splash, arkrt, display daemon)");
  {
    /* Check each binary individually for incremental builds */
    char arkrt_out[512], daemon_out[512];
    snprintf(arkrt_out, sizeof(arkrt_out), "%s/arkrt", staging);
    snprintf(daemon_out, sizeof(daemon_out), "%s/system/ui_daemon", staging);

    /* arkrt sources */
    char a1[512], a2[512], a3[512], a4[512], a5[512], a6[512], a7[512];
    snprintf(a1, sizeof(a1), "%s/arkrt/main.swift", base);
    snprintf(a2, sizeof(a2), "%s/arkrt/KernelBridge.swift", base);
    snprintf(a3, sizeof(a3), "%s/arkrt/CommandRouter.swift", base);
    snprintf(a4, sizeof(a4), "%s/arkrt/IPC.swift", base);
    snprintf(a5, sizeof(a5), "%s/arkrt/ServiceManager.swift", base);
    snprintf(a6, sizeof(a6), "%s/arkrt/NetworkService.swift", base);
    snprintf(a7, sizeof(a7), "%s/arkrt/InputService.swift", base);
    const char *arkrt_sources[] = {a1, a2, a3, a4, a5, a6, a7};

    if (is_up_to_date(arkrt_out, arkrt_sources, 7)) {
      printf("  \033[0;32m✓\033[0m arkrt is up-to-date (skipping)\n");
    } else {
      char cmd[1024];
      snprintf(cmd, sizeof(cmd), "make -C %s %s/arkrt", base,
               is_rpi4 ? "out_staging/rpi4" : "out_staging");
      run(cmd);
    }

    /* ark_compositor sources */
    char c1[512], c2[512], c3[512], c4[512];
    snprintf(c1, sizeof(c1), "%s/arkrt/compositor/main.swift", base);
    snprintf(c2, sizeof(c2), "%s/arkrt/compositor/Compositor.swift", base);
    snprintf(c3, sizeof(c3), "%s/arkrt/compositor/ClientProtocol.swift", base);
    snprintf(c4, sizeof(c4), "%s/arkrt/compositor/drm_helper.c", base);
    const char *comp_sources[] = {c1, c2, c3, c4};
    char comp_out[512];
    snprintf(comp_out, sizeof(comp_out), "%s/system/ark_compositor", staging);

    if (is_up_to_date(comp_out, comp_sources, 4)) {
      printf("  \033[0;32m✓\033[0m ark_compositor is up-to-date (skipping)\n");
    } else {
      char cmd[1024];
    }
  }

  /* ─── Step 10: Verify ArkGraphics framework ────────────────── */
  step(is_rpi4 ? "Verifying graphics & input components (ARM64)"
               : "Verifying graphics & input components");
  {
    printf("  \033[0;32m✓\033[0m ArkGraphics & ArkInput ready at "
           "arkrt/ArkGraphics\n");
  }

  /* ─── Step 11: Compile setup_app (incremental) ─────────────── */
  step(is_rpi4 ? "Cross-compiling setup_app (ARM64)"
               : "Compiling setup_app (setup app)");
  {
    char ut_out[512];
    snprintf(ut_out, sizeof(ut_out), "%s/system/setup_app", staging);

    char u1[512], u2[512], u3[512], u4[512], u5[512], u6[512], u7[512], u8[512];
    snprintf(u1, sizeof(u1), "%s/system/apps/setup_app.swift", base);
    snprintf(u2, sizeof(u2), "%s/arkrt/ArkGraphics/ArkGraphics.swift", base);
    snprintf(u3, sizeof(u3), "%s/arkrt/ArkGraphics/ArkWrite.swift", base);
    snprintf(u4, sizeof(u4), "%s/arkrt/ArkGraphics/ArkShapes.swift", base);
    snprintf(u5, sizeof(u5), "%s/arkrt/ArkGraphics/ArkFontRobotoBold.swift",
             base);
    snprintf(u6, sizeof(u6), "%s/arkrt/ArkGraphics/ArkInput.swift", base);
    snprintf(u7, sizeof(u7), "%s/arkrt/ArkGraphics/ArkUI.swift", base);
    snprintf(u8, sizeof(u8), "%s/arkrt/ArkGraphics/ArkIcons.swift", base);
    const char *ut_sources[] = {u1, u2, u3, u4, u5, u6, u7, u8};

    if (is_up_to_date(ut_out, ut_sources, 8)) {
      printf("  \033[0;32m✓\033[0m setup_app is up-to-date (skipping)\n");
    } else {
      char cmd[1024];
      snprintf(cmd, sizeof(cmd), "make -C %s %s/system/setup_app", base,
               is_rpi4 ? "out_staging/rpi4" : "out_staging");
      run(cmd);
    }
    
    char ui_cmd[1024];
    snprintf(ui_cmd, sizeof(ui_cmd), "make -C %s %s/system/ui_daemon", base,
             is_rpi4 ? "out_staging/rpi4" : "out_staging");
    run(ui_cmd);
  }

  /* ─── Step 12: Sign arkrt (Verified Boot) ──────────────────── */
  step(is_rpi4 ? "Signing arkrt binary (Verified Boot, ARM64)"
               : "Signing arkrt binary (Verified Boot)");
  {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd),
             "python3 %s/verify/sign.py %s/verify/securebuild.ark "
             "%s/arkrt %s/signature.bin",
             vendor_dir, vendor_dir, staging, staging);
    run(cmd);
  }

  /* ─── Step 13: (Skipped) init (PID 1) is now handled natively by arkrt ─── */

  /* ─── Step 14: Pack initramfs ──────────────────────────────── */
  step("Packing initramfs (base + runtime + display + UI)");
  {
    char *initramfs_ext = pjoin(staging, "initramfs_ext");
    char *base_initramfs = pjoin(boot_dir, "initramfs.img");

    char cmd[2048];

    /* Extract base initramfs (contains busybox, kernel modules, etc.) if
     * present */
    snprintf(cmd, sizeof(cmd), "rm -rf %s && mkdir -p %s", initramfs_ext,
             initramfs_ext);
    run(cmd);

    if (access(base_initramfs, F_OK) != -1) {
      snprintf(cmd, sizeof(cmd),
               "cd %s && zcat %s 2>/dev/null | cpio -id --no-preserve-owner "
               "2>/dev/null || true",
               initramfs_ext, base_initramfs);
      run(cmd);
    }

    /* Create required directories in initramfs */
    snprintf(cmd, sizeof(cmd),
             "mkdir -p %s/system/services %s/run %s/tmp %s/var/log",
             initramfs_ext, initramfs_ext, initramfs_ext, initramfs_ext);
    run(cmd);

    /* Copy arkrt directly as /init (PID 1) */
    snprintf(cmd, sizeof(cmd), "cp %s/arkrt %s/init && chmod +x %s/init",
             staging, initramfs_ext, initramfs_ext);
    run(cmd);

    snprintf(cmd, sizeof(cmd), "cp %s/signature.bin %s/signature.bin", staging,
             initramfs_ext);
    run(cmd);

    /* Copy display daemon and setup app into /system/ */
    snprintf(cmd, sizeof(cmd),
             "cp %s/system/ui_daemon %s/system/ui_daemon && "
             "chmod +x %s/system/ui_daemon 2>/dev/null || true",
             staging, initramfs_ext, initramfs_ext);
    run(cmd);

    snprintf(cmd, sizeof(cmd),
             "cp %s/system/setup_app %s/system/setup_app && "
             "chmod +x %s/system/setup_app 2>/dev/null || true",
             staging, initramfs_ext, initramfs_ext);
    run(cmd);

    /* Copy service definition files */
    snprintf(cmd, sizeof(cmd),
             "cp -r %s/services/* %s/system/services/ 2>/dev/null || true",
             staging_sys, initramfs_ext);
    run(cmd);

    /* Repack initramfs as gzip-compressed cpio archive */
    snprintf(cmd, sizeof(cmd),
             "cd %s && find . | cpio -H newc -o 2>/dev/null | gzip > "
             "%s/initramfs.img",
             initramfs_ext, out_dir);
    run(cmd);

    free(initramfs_ext);
    free(base_initramfs);
  }

  /* ─── Step 15: Assemble bootloader + stage2 ────────────────── */
  step("Assembling bootloader");
  {
    char cmd[1024];

    /* Generate boot animation frames from Python script */
    snprintf(cmd, sizeof(cmd),
             "python3 %s/source/animationframes/generate_frames.py %s",
             boot_dir, staging);
    run(cmd);

    snprintf(cmd, sizeof(cmd),
             "make -C %s out_staging/bootloader.bin out_staging/stage2.bin "
             "out_staging/BOOTX64.EFI",
             base);
    run(cmd);
  }

  /* ─── Step 16: Generate disk images ────────────────────────── */
  step(is_rpi4 ? "Generating disk images (RPi4 SD Card)"
               : "Generating disk images");
  {
    char cmd[2048];

    /* boot.img — bootloader + kernel + initramfs packed by pack_boot.py */
    if (!is_rpi4) {
      snprintf(cmd, sizeof(cmd), "python3 %s/tools/pack_boot.py --base-dir %s",
               repo, base);
      run(cmd);
    } else {
      /* Copy ARM64 kernel to finished/rpi4 as kernel8.img */
      snprintf(cmd, sizeof(cmd), "cp %s/prebuilts/arm64 %s/kernel8.img",
               kernel_dir, out_dir);
      run(cmd);
      snprintf(cmd, sizeof(cmd),
               "cp %s/prebuilts/*.dtb %s/ 2>/dev/null || true", kernel_dir,
               out_dir);
      run(cmd);
      snprintf(cmd, sizeof(cmd),
               "mkdir -p %s/overlays && cp %s/prebuilts/rpi4_dtbofiles/*.dtbo "
               "%s/overlays/ 2>/dev/null || true",
               out_dir, kernel_dir, out_dir);
      run(cmd);
    }

    /* sys.img — system partition (frameworks, apps, services) */
    snprintf(cmd, sizeof(cmd),
             "dd if=/dev/zero of=%s/sys.img bs=1M count=2048 2>/dev/null",
             out_dir);
    run(cmd);
    snprintf(cmd, sizeof(cmd), "mkfs.ext4 -d %s %s/sys.img 2>/dev/null",
             staging_sys, out_dir);
    run(cmd);

    /* vend.img — vendor partition (OEM blobs, keys) */
    snprintf(cmd, sizeof(cmd),
             "dd if=/dev/zero of=%s/vend.img bs=1M count=100 2>/dev/null",
             out_dir);
    run(cmd);
    snprintf(cmd, sizeof(cmd),
             "mkfs.ext4 -F -d %s %s/vend.img 2>/dev/null || mkfs.ext4 -F "
             "%s/vend.img 2>/dev/null",
             staging_vend, out_dir, out_dir);
    run(cmd);

    /* ── dtbo.img — Device Tree Blobs + Overlays ────────────── */
    /* Stage all .dtb files from kernel/prebuilts/ and all .dtbo files
     * from kernel/prebuilts/rpi4_dtbofiles/ into a temp directory,
     * then pack them into an ext4 image. The bootloader extracts
     * these and feeds them to the kernel. */
    {
      char *dtbo_stage = pjoin(staging, "dtbo_staging");
      snprintf(cmd, sizeof(cmd), "rm -rf %s && mkdir -p %s/overlays",
               dtbo_stage, dtbo_stage);
      run(cmd);
      /* Copy .dtb files (base device trees) */
      snprintf(cmd, sizeof(cmd),
               "cp %s/prebuilts/*.dtb %s/ 2>/dev/null || true", kernel_dir,
               dtbo_stage);
      run(cmd);
      /* Copy .dtbo files (overlays) */
      snprintf(cmd, sizeof(cmd),
               "cp %s/prebuilts/rpi4_dtbofiles/*.dtbo %s/overlays/ 2>/dev/null "
               "|| true",
               kernel_dir, dtbo_stage);
      run(cmd);
      /* Also copy any .dtb files that are in the rpi4_dtbofiles dir */
      snprintf(cmd, sizeof(cmd),
               "cp %s/prebuilts/rpi4_dtbofiles/*.dtb %s/overlays/ 2>/dev/null "
               "|| true",
               kernel_dir, dtbo_stage);
      run(cmd);
      /* Calculate size: round up to nearest MB + 2MB headroom */
      snprintf(cmd, sizeof(cmd),
               "du -sm %s | awk '{print ($1 < 4 ? 4 : $1 + 2)}'", dtbo_stage);
      FILE *p = popen(cmd, "r");
      int dtbo_size_mb = 4; /* default fallback */
      if (p) {
        fscanf(p, "%d", &dtbo_size_mb);
        pclose(p);
      }
      printf("  dtbo.img size: %d MB\n", dtbo_size_mb);
      snprintf(cmd, sizeof(cmd),
               "dd if=/dev/zero of=%s/dtbo.img bs=1M count=%d 2>/dev/null",
               out_dir, dtbo_size_mb);
      run(cmd);
      snprintf(cmd, sizeof(cmd), "mkfs.ext4 -d %s %s/dtbo.img 2>/dev/null",
               dtbo_stage, out_dir);
      run(cmd);
      printf("  \033[0;32m✓\033[0m dtbo.img packed with device tree blobs and "
             "overlays\n");
      free(dtbo_stage);
    }

    /* ── vbk.img — Verified Boot Key ───────────────────────────── */
    /* Contains the securebuild.ark signing key. The bootloader reads
     * the VBK key from here and verifies that sys.img and vend.img
     * were signed with the matching key. */
    {
      char *securebuild_path = pjoin(vendor_dir, "verify/securebuild.ark");
      ArkConfig *sbc = ark_parse(securebuild_path);
      int has_keys = 0;
      for (int i = 0; i < sbc->standalone_count; i++) {
        if (strstr(sbc->standalone[i], "ARK-OS-") != NULL) {
          has_keys = 1;
          break;
        }
      }
      if (has_keys) {
        printf("  \033[0;32m✓\033[0m VBK signing keys found — generating "
               "vbk.img\n");
        char *vbk_stage = pjoin(staging, "vbk_staging");
        snprintf(cmd, sizeof(cmd), "rm -rf %s && mkdir -p %s", vbk_stage,
                 vbk_stage);
        run(cmd);
        snprintf(cmd, sizeof(cmd), "cp %s %s/securebuild.ark", securebuild_path,
                 vbk_stage);
        run(cmd);
        snprintf(cmd, sizeof(cmd),
                 "dd if=/dev/zero of=%s/vbk.img bs=1M count=1 2>/dev/null",
                 out_dir);
        run(cmd);
        snprintf(cmd, sizeof(cmd), "mkfs.ext4 -d %s %s/vbk.img 2>/dev/null",
                 vbk_stage, out_dir);
        run(cmd);
        printf("  \033[0;32m✓\033[0m vbk.img packed with Verified Boot Key\n");
        free(vbk_stage);
      } else {
        printf("  \033[1;33m⚠\033[0m No signing keys — skipping vbk.img\n");
      }
      ark_free(sbc);
      free(securebuild_path);
    }

    /* ── vbmeta.img — Boot Metadata / Mount Configuration ──── */
    /* Contains vbmeta.ark which tells the bootloader the partition
     * layout and how to mount each image. Loaded first during boot. */
    {
      char *vbmeta_stage = pjoin(staging, "vbmeta_staging");
      char *vbmeta_ark = pjoin(boot_dir, "vbmeta.ark");
      snprintf(cmd, sizeof(cmd), "rm -rf %s && mkdir -p %s", vbmeta_stage,
               vbmeta_stage);
      run(cmd);
      if (access(vbmeta_ark, F_OK) != -1) {
        snprintf(cmd, sizeof(cmd), "cp %s %s/vbmeta.ark", vbmeta_ark,
                 vbmeta_stage);
        run(cmd);
      } else {
        printf("  \033[1;33m⚠\033[0m boot/vbmeta.ark not found — vbmeta.img "
               "will be empty\n");
      }
      snprintf(cmd, sizeof(cmd),
               "dd if=/dev/zero of=%s/vbmeta.img bs=1M count=1 2>/dev/null",
               out_dir);
      run(cmd);
      snprintf(cmd, sizeof(cmd), "mkfs.ext4 -d %s %s/vbmeta.img 2>/dev/null",
               vbmeta_stage, out_dir);
      run(cmd);
      printf(
          "  \033[0;32m✓\033[0m vbmeta.img packed with mount configuration\n");
      free(vbmeta_stage);
      free(vbmeta_ark);
    }

    /* ── rpi4.img — Assemble all images into one flashable SD card image ── */
    if (is_rpi4 && !no_image) {
      printf("\n  Assembling rpi4.img (flashable SD card image)...\n");
      snprintf(cmd, sizeof(cmd), "python3 %s/tools/pack_rpi4.py --base-dir %s",
               repo, base);
      run(cmd);
    } else if (is_rpi4 && no_image) {
      printf("\n  \033[1;33m[--no-image]\033[0m Skipping rpi4.img assembly — "
             "individual .img files available in %s/\n",
             out_dir);
    }
  }

  printf("\n\033[1;32m╔══════════════════════════════════════╗\033[0m\n");
  printf("\033[1;32m║        Build Complete!                ║\033[0m\n");
  printf("\033[1;32m╚══════════════════════════════════════╝\033[0m\n");
  printf("Run \033[1mmake test-uefi or make test-bios (if x86_64) \033[0m or "
         "\033[1mmake test-uefi-arm64 or make test-bios-arm64 (if "
         "aarch64)\033[0m to launch ArkOS in QEMU.\n\n");

  /* Free all heap-allocated path strings */
  free(repo);
  free(out_dir);
  free(staging);
  free(staging_sys);
  free(staging_fw);
  free(staging_vend);
  free(staging_boot);
  free(fw_dir);
  free(boot_dir);
  free(vendor_dir);
  free(system_dir);
  free(kernel_dir);

  return 0;
}