#!/usr/bin/env zsh
set -e

ARCH=${1:-riscv64}

# Architecture-specific settings
if [ "$ARCH" = "aarch64" ]; then
  NIM_CPU="arm64"
  ZIG_TARGET="aarch64-freestanding-none"
  ZIG_CPU="baseline"
  LINKER_SCRIPT="apps/linker_user_aarch64.ld"
  BUILD_FLAG="-Darch=aarch64"
  echo "=== Building NipBox Userland (aarch64) ==="
else
  NIM_CPU="riscv64"
  ZIG_TARGET="riscv64-freestanding-none"
  ZIG_CPU="sifive_u54"
  LINKER_SCRIPT="apps/linker_user.ld"
  BUILD_FLAG=""
  echo "=== Building NipBox Userland (riscv64) ==="
fi

# Compile Nim sources to C
nim c --cpu:${NIM_CPU} --os:any --compileOnly --mm:arc --opt:size \
  --stackTrace:off --lineDir:off --nomain --nimcache:build/init_nimcache \
  -d:noSignalHandler -d:RUMPK_USER -d:nimAllocPagesViaMalloc -d:NIPBOX_LITE \
  npl/nipbox/nipbox.nim

# Compile Nim-generated C (check if files exist first)
# Skip net_glue (needs LwIP headers not available in userland build)
EXTRA_CC_FLAGS=""
if [ "$ARCH" = "riscv64" ]; then
  EXTRA_CC_FLAGS="-mcmodel=medany"
fi

if ls build/init_nimcache/*.c 1> /dev/null 2>&1; then
  for f in build/init_nimcache/*.c; do
    case "$f" in
      *net_glue*) echo "  [skip] $f (LwIP dependency)"; continue ;;
    esac
    zig cc -target ${ZIG_TARGET} -mcpu=${ZIG_CPU} ${EXTRA_CC_FLAGS} \
      -fno-sanitize=all -fno-vectorize \
      -I/usr/lib/nim/lib -Icore -Ilibs/membrane -Ilibs/membrane/include \
      -include string.h \
      -Os -c "$f" -o "${f%.c}.o"
  done
fi

# Compile clib
zig cc -target ${ZIG_TARGET} -mcpu=${ZIG_CPU} ${EXTRA_CC_FLAGS} \
  -fno-sanitize=all \
  -DNO_SYS=1 -DOMIT_EXIT -DRUMPK_USER -Ilibs/membrane/include -c libs/membrane/clib.c -o build/clib_user.o

# Create startup assembly
if [ "$ARCH" = "aarch64" ]; then
  cat > build/head_user.S << 'EOF'
.section .text._start
.global _start
_start:
    bl NimMain
1:  wfi
    b 1b
EOF
else
  cat > build/head_user.S << 'EOF'
.section .text._start
.global _start
_start:
    .option push
    .option norelax
    1:auipc gp, %pcrel_hi(__global_pointer$)
      addi  gp, gp, %pcrel_lo(1b)
    .option pop
    call NimMain
    1: wfi
       j 1b
EOF
fi

zig cc -target ${ZIG_TARGET} -mcpu=${ZIG_CPU} ${EXTRA_CC_FLAGS} \
  -fno-sanitize=all \
  -c build/head_user.S -o build/head_user.o

# Link init
zig cc -target ${ZIG_TARGET} -mcpu=${ZIG_CPU} ${EXTRA_CC_FLAGS} -nostdlib \
  -fno-sanitize=all \
  -T ${LINKER_SCRIPT} -Wl,--gc-sections \
  build/head_user.o build/init_nimcache/*.o build/clib_user.o \
  -o build/init

echo "✓ NipBox binary built (${ARCH})"
file build/init

# Create initrd
mkdir -p build/sysro/bin
cp build/init build/sysro/init
if [ "$ARCH" = "riscv64" ] && [ -f vendor/mksh/mksh.elf ]; then
  cp vendor/mksh/mksh.elf build/sysro/bin/mksh
fi
cd build/sysro
tar --format=ustar -cf ../initrd.tar *
cd ../..

# Embed initrd
cat > build/embed_initrd.S << EOF
.section .rodata
.global _initrd_start
.global _initrd_end
.align 4
_initrd_start:
    .incbin "$(pwd)/build/initrd.tar"
_initrd_end:
EOF

zig cc -target ${ZIG_TARGET} -mcpu=${ZIG_CPU} ${EXTRA_CC_FLAGS} \
  -c build/embed_initrd.S -o build/initrd.o

cp build/initrd.tar hal/initrd.tar

# Build kernel
rm -f zig-out/lib/librumpk_hal.a
zig build ${BUILD_FLAG}

echo "=== BUILD COMPLETE (${ARCH}) ==="
ls -lh build/init zig-out/bin/rumpk.elf