/*
 * ArkOS Build System
 * ------------------
 * Orchestrates the full OS build pipeline.
 * All paths are relative to the base directory passed as argv[1].
 *
 * 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 step counter ─────────────────────────────────────────── */

#define TOTAL_STEPS 14

static int current_step = 0;

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

/* ── Helpers ────────────────────────────────────────────────────── */

static void 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);
  }
}

/* Join two path components. Caller must free. */
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;
}

/* ── Main ───────────────────────────────────────────────────────── */

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

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

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

  /* Key directories */
  char *out_dir      = pjoin(base, "finished");
  char *staging      = 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");

  /* Compiler for OS components (host gcc for now) */
  const char *cc = "gcc";

  /* 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;35mArkOS Build System v2.0\033[0m\n");

  /* ─── Step 1: Clean ─────────────────────────────────────────── */
  step("Cleaning previous build");
  {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd),
             "chmod -R 777 %s 2>/dev/null || true", staging);
    run(cmd);
    snprintf(cmd, sizeof(cmd), "rm -rf %s %s", out_dir, staging);
    run(cmd);
  }

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

  /* ─── Step 3: Copy frameworks ───────────────────────────────── */
  step("Integrating frameworks");
  {
    DIR *d = opendir(fw_dir);
    if (d) {
      struct dirent *ent;
      while ((ent = readdir(d)) != NULL) {
        if (ent->d_name[0] == '.')
          continue;
        /* Skip excluded frameworks */
        if (strcmp(ent->d_name, "DRM") == 0 ||
            strcmp(ent->d_name, "SwiftUIFramework") == 0) {
          printf("  Skipping excluded framework: %s\n", ent->d_name);
          continue;
        }
        char cmd[1024];
        snprintf(cmd, sizeof(cmd), "cp -r %s/%s %s/",
                 fw_dir, ent->d_name, staging_fw);
        run(cmd);
      }
      closedir(d);
    }
  }

  /* ─── Step 4: Copy boot files ───────────────────────────────── */
  step("Preparing boot files");
  {
    DIR *d = opendir(boot_dir);
    if (d) {
      struct dirent *ent;
      while ((ent = readdir(d)) != NULL) {
        if (ent->d_name[0] == '.')
          continue;
        if (strcmp(ent->d_name, "DRM") == 0) {
          printf("  Skipping boot/DRM\n");
          continue;
        }
        char cmd[1024];
        snprintf(cmd, sizeof(cmd), "cp -r %s/%s %s/",
                 boot_dir, ent->d_name, staging_boot);
        run(cmd);
      }
      closedir(d);
    }
  }

  /* ─── Step 5: Copy vendor files ─────────────────────────────── */
  step("Preparing vendor files");
  {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd),
             "cp -r %s/* %s/ 2>/dev/null || true",
             vendor_dir, staging_vend);
    run(cmd);
  }

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

  /* ─── Step 7: Compile fb_helper.c ───────────────────────────── */
  step("Compiling : system/fb_helper.c");
  {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd), "make -C %s out_staging/fb_helper.o", base);
    run(cmd);
  }

  /* ─── Step 8: Compile swift_splash.swift & arkrt ────────────── */
  step("Compiling : system/swift_splash.swift & arkrt");
  {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd), "make -C %s out_staging/swift_splash out_staging/arkrt", base);
    run(cmd);
  }

  /* ─── Step 9: Sign arkrt (Verified Boot) ────────────────────── */
  step("Signing arkrt (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 10: Compile init.c + vendor/verify/sha256.c ──── */
  step("Compiling : system/init.c + vendor/verify/sha256.c");
  {
    char cmd2[1024];
    snprintf(cmd2, sizeof(cmd2), "make -C %s out_staging/init", base);
    run(cmd2);
  }

  /* ─── Step 11: Pack initramfs ───────────────────────────────── */
  step("Packing initramfs");
  {
    char *initramfs_ext = pjoin(staging, "initramfs_ext");
    char *base_initramfs = pjoin(boot_dir, "initramfs.img");

    char cmd[2048];

    /* Extract base initramfs */
    snprintf(cmd, sizeof(cmd),
             "rm -rf %s && mkdir -p %s", initramfs_ext, initramfs_ext);
    run(cmd);

    snprintf(cmd, sizeof(cmd),
             "cd %s && zcat %s | cpio -id 2>/dev/null",
             initramfs_ext, base_initramfs);
    run(cmd);

    /* Copy compiled binaries from staging into initramfs */
    snprintf(cmd, sizeof(cmd),
             "cp %s/init %s/init && chmod +x %s/init",
             staging, initramfs_ext, initramfs_ext);
    run(cmd);

    snprintf(cmd, sizeof(cmd),
             "cp %s/swift_splash %s/swift_splash && chmod +x %s/swift_splash",
             staging, initramfs_ext, initramfs_ext);
    run(cmd);

    snprintf(cmd, sizeof(cmd),
             "cp %s/arkrt %s/arkrt && chmod +x %s/arkrt",
             staging, initramfs_ext, initramfs_ext);
    run(cmd);

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

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

    free(initramfs_ext);
    free(base_initramfs);
  }

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

    /* Generate animation frames */
    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 13: Generate boot.img ────────────────────────────── */
  step("Generating boot.img");
  {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd),
             "python3 %s/tools/pack_boot.py --base-dir %s",
             repo, base);
    run(cmd);
  }

  /* ─── Step 14: Generate disk images ─────────────────────────── */
  step("Generating disk images");
  {
    char cmd[1024];

    /* sys.img */
    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 */
    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 -d %s %s/vend.img 2>/dev/null",
             staging_vend, out_dir);
    run(cmd);

    /* avb.img (Verified Boot) */
    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 Signing keys found, generating avb.img\n");
      snprintf(cmd, sizeof(cmd),
               "dd if=/dev/zero of=%s/avb.img bs=1M count=1 2>/dev/null",
               out_dir);
      run(cmd);
    } else {
      printf("  \033[1;33m⚠\033[0m No signing keys found, skipping avb.img\n");
    }
    ark_free(sbc);
    free(securebuild_path);

    /* vbmeta.img + dtbo.img */
    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),
             "dd if=/dev/zero of=%s/dtbo.img bs=1M count=1 2>/dev/null",
             out_dir);
    run(cmd);
  }

  printf("\n\033[1;32mBuild complete!\033[0m Run \033[1mmake run\033[0m to test.\n\n");

  /* Cleanup heap */
  free(repo);
  free(out_dir);
  free(staging);
  free(staging_sys);
  free(staging_fw);
  free(staging_vend);
  free(staging_boot);
  /* cc is a string literal, no free needed */
  free(fw_dir);
  free(boot_dir);
  free(vendor_dir);
  free(system_dir);
  free(kernel_dir);

  return 0;
}