#!/usr/bin/env bash
set -Eeuo pipefail
umask 022

SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)
PROJECT_DIR=$(cd -- "${SCRIPT_DIR}/.." && pwd)

BASE_XZ="${PROJECT_DIR}/downloads/Armbian_community_26.11.0-trunk.1_Turing-rk1_trixie_vendor_6.1.115_minimal.img.xz"
BASE_SHA256="cbbd16d90786331fbfbbb695614a53f5e87ea441ea37b425d986334318190cf6"
KODI_DEB="${PROJECT_DIR}/media/out/rk1-media-stack_20260721.1741+rk1.1_arm64.deb"
DEB_DIR="${PROJECT_DIR}/work/apt/archives"
RUNTIME_DIR="${PROJECT_DIR}/runtime"
RKNN_SOURCE_DIR="${PROJECT_DIR}/downloads/rknn-toolkit2-v2.3.2-minimal"
SSH_PUBLIC_KEY=""
OUTPUT_DIR="${PROJECT_DIR}/dist"
IMAGE_SIZE_MIB=4096
IMAGE_REVISION="r4"

usage() {
  cat <<'EOF'
Usage: build-image.sh --ssh-public-key PATH [options]

Options:
  --base-xz PATH          Verified Armbian .img.xz base
  --kodi-deb PATH         Repacked RKMPP Kodi arm64 package
  --deb-dir PATH          Offline Debian dependency directory
  --runtime-dir PATH      RKNN/self-test runtime project
  --rknn-source-dir PATH  Offline RKNN Toolkit2 v2.3.2 asset bundle
  --output-dir PATH       Artifact directory
  --image-size-mib N      Expanded image size (default: 4096)
  --revision NAME         Artifact revision suffix (default: r4)
EOF
}

while (($#)); do
  case "$1" in
    --ssh-public-key) SSH_PUBLIC_KEY=${2:?}; shift 2 ;;
    --base-xz) BASE_XZ=${2:?}; shift 2 ;;
    --kodi-deb) KODI_DEB=${2:?}; shift 2 ;;
    --deb-dir) DEB_DIR=${2:?}; shift 2 ;;
    --runtime-dir) RUNTIME_DIR=${2:?}; shift 2 ;;
    --rknn-source-dir) RKNN_SOURCE_DIR=${2:?}; shift 2 ;;
    --output-dir) OUTPUT_DIR=${2:?}; shift 2 ;;
    --image-size-mib) IMAGE_SIZE_MIB=${2:?}; shift 2 ;;
    --revision) IMAGE_REVISION=${2:?}; shift 2 ;;
    -h|--help) usage; exit 0 ;;
    *) echo "Unknown argument: $1" >&2; usage >&2; exit 2 ;;
  esac
done

for command_name in debugfs e2fsck install jq resize2fs sfdisk sgdisk sha256sum ssh-keygen xz; do
  command -v "$command_name" >/dev/null || { echo "Missing command: $command_name" >&2; exit 1; }
done
[[ -n "$SSH_PUBLIC_KEY" && -s "$SSH_PUBLIC_KEY" ]] || {
  echo "--ssh-public-key must name a non-empty public-key file" >&2
  exit 2
}
grep -Eq -- '-----BEGIN .*PRIVATE KEY-----' "$SSH_PUBLIC_KEY" && {
  echo "Refusing a private key: $SSH_PUBLIC_KEY" >&2
  exit 2
}
awk '
  /^[[:space:]]*($|#)/ { next }
  $1 !~ /^(ssh-ed25519|ssh-rsa|ecdsa-sha2-nistp(256|384|521)|sk-ssh-ed25519@openssh.com|sk-ecdsa-sha2-nistp256@openssh.com)$/ { exit 1 }
  NF < 2 { exit 1 }
  { count++ }
  END { if (count < 1) exit 1 }
' "$SSH_PUBLIC_KEY" || {
  echo "Use plain OpenSSH public-key lines without authorized_keys options" >&2
  exit 2
}
ssh-keygen -lf "$SSH_PUBLIC_KEY" >/dev/null || {
  echo "ssh-keygen rejected the public-key file" >&2
  exit 2
}
[[ -s "$BASE_XZ" ]] || { echo "Missing base image: $BASE_XZ" >&2; exit 1; }
[[ -s "$KODI_DEB" ]] || { echo "Missing Kodi package: $KODI_DEB" >&2; exit 1; }
[[ -d "$DEB_DIR" ]] || { echo "Missing dependency directory: $DEB_DIR" >&2; exit 1; }
[[ -x "$RUNTIME_DIR/install-rknn-runtime.sh" ]] || {
  echo "Missing RKNN installer below: $RUNTIME_DIR" >&2
  exit 1
}
[[ -d "$RKNN_SOURCE_DIR" ]] || {
  echo "Missing offline RKNN source bundle: $RKNN_SOURCE_DIR" >&2
  exit 1
}
[[ "$IMAGE_SIZE_MIB" =~ ^[0-9]+$ ]] && ((IMAGE_SIZE_MIB >= 3072)) || {
  echo "Image size must be an integer of at least 3072 MiB" >&2
  exit 2
}
[[ "$IMAGE_REVISION" =~ ^[a-z0-9][a-z0-9._-]*$ ]] || {
  echo "Image revision contains unsafe characters: $IMAGE_REVISION" >&2
  exit 2
}

actual_base_sha=$(sha256sum "$BASE_XZ" | awk '{print $1}')
[[ "$actual_base_sha" == "$BASE_SHA256" ]] || {
  echo "Base image checksum mismatch: $actual_base_sha" >&2
  exit 1
}

repair_ext4() {
  local status
  if e2fsck -fy "$1" >/dev/null; then
    return 0
  else
    status=$?
  fi
  # fsck bit 0 means errors were corrected; this is expected after debugfs.
  ((status == 1)) || {
    echo "e2fsck failed for $1 with status $status" >&2
    return "$status"
  }
}

BUILD_ID="$(date -u +%Y%m%d)-${IMAGE_REVISION}-trixie-vendor-6.1.115"
WORK_DIR="${PROJECT_DIR}/work/image-${BUILD_ID}"
STAGE_DIR="${WORK_DIR}/overlay"
RAW_IMAGE="${WORK_DIR}/rk1-media-${BUILD_ID}.img"
ROOTFS_IMAGE="${WORK_DIR}/rootfs.ext4"
OUTPUT_BASENAME="rk1-media-${BUILD_ID}-turing-rk1-emmc"

rm -rf -- "$WORK_DIR"
mkdir -p "$STAGE_DIR" "$OUTPUT_DIR"
cp -a "${SCRIPT_DIR}/overlay/." "$STAGE_DIR/"
mkdir -p "$STAGE_DIR/opt/rk1-seed/debs" "$STAGE_DIR/usr/share/doc/rk1-media-image"
install -D -m 0755 "$PROJECT_DIR/scripts/live-migrate-root-to-nvme.sh" \
  "$STAGE_DIR/usr/local/sbin/rk1-media-migrate-root-to-nvme"
cp -a "$DEB_DIR"/*.deb "$STAGE_DIR/opt/rk1-seed/debs/"
cp -a "$KODI_DEB" "$STAGE_DIR/opt/rk1-seed/debs/"
cp -a "$SSH_PUBLIC_KEY" "$STAGE_DIR/opt/rk1-seed/authorized_keys"
(
  cd "$STAGE_DIR/opt/rk1-seed/debs"
  sha256sum -- *.deb >../debs.sha256
)
cp -a "$RUNTIME_DIR/README.md" "$RUNTIME_DIR/rknn-version.env" \
  "$STAGE_DIR/usr/share/doc/rk1-media-image/"
"$RUNTIME_DIR/install-rknn-runtime.sh" \
  --rootfs "$STAGE_DIR" --source-dir "$RKNN_SOURCE_DIR"

xz --decompress --stdout "$BASE_XZ" >"$RAW_IMAGE"
partition_json=$(sfdisk --json "$RAW_IMAGE")
partition_start=$(jq -r '.partitiontable.partitions[0].start' <<<"$partition_json")
partition_size=$(jq -r '.partitiontable.partitions[0].size' <<<"$partition_json")
partition_type=$(jq -r '.partitiontable.partitions[0].type' <<<"$partition_json")
partition_uuid=$(jq -r '.partitiontable.partitions[0].uuid' <<<"$partition_json")
[[ "$partition_start" == "32768" ]] || {
  echo "Unexpected base partition start: $partition_start" >&2
  exit 1
}

dd if="$RAW_IMAGE" of="$ROOTFS_IMAGE" bs=512 skip="$partition_start" count="$partition_size" status=none
repair_ext4 "$ROOTFS_IMAGE"

firstrun_file="${WORK_DIR}/armbian-firstrun"
debugfs -R 'cat /etc/default/armbian-firstrun' "$ROOTFS_IMAGE" \
  2>/dev/null >"$firstrun_file"
if grep -q '^OPENSSHD_REGENERATE_HOST_KEYS=' "$firstrun_file"; then
  sed -i 's/^OPENSSHD_REGENERATE_HOST_KEYS=.*/OPENSSHD_REGENERATE_HOST_KEYS=false/' \
    "$firstrun_file"
else
  printf '\nOPENSSHD_REGENERATE_HOST_KEYS=false\n' >>"$firstrun_file"
fi
mkdir -p "$STAGE_DIR/etc/default"
cp "$firstrun_file" "$STAGE_DIR/etc/default/armbian-firstrun"

# Remove the generic base image's root password hash before the image exists.
shadow_file="${WORK_DIR}/shadow"
debugfs -R 'cat /etc/shadow' "$ROOTFS_IMAGE" 2>/dev/null \
  | awk 'BEGIN { FS=OFS=":" } $1 == "root" { $2="!"; found=1 } { print } END { if (!found) exit 1 }' \
  >"$shadow_file"
mkdir -p "$STAGE_DIR/etc"
cp "$shadow_file" "$STAGE_DIR/etc/shadow"

truncate -s "${IMAGE_SIZE_MIB}M" "$RAW_IMAGE"
sgdisk --move-second-header "$RAW_IMAGE" >/dev/null
sgdisk --delete=1 "$RAW_IMAGE" >/dev/null
sgdisk --new="1:${partition_start}:0" --typecode="1:${partition_type}" --change-name='1:rootfs' "$RAW_IMAGE" >/dev/null
sgdisk --partition-guid="1:${partition_uuid}" "$RAW_IMAGE" >/dev/null

partition_json=$(sfdisk --json "$RAW_IMAGE")
new_partition_size=$(jq -r '.partitiontable.partitions[0].size' <<<"$partition_json")
truncate -s "$((new_partition_size * 512))" "$ROOTFS_IMAGE"
resize2fs "$ROOTFS_IMAGE" >/dev/null

env_file="${WORK_DIR}/armbianEnv.txt"
debugfs -R 'cat /boot/armbianEnv.txt' "$ROOTFS_IMAGE" 2>/dev/null >"$env_file"
if sed -n 's/^overlays=//p' "$env_file" | tr '[:space:]' '\n' | grep -Fxq panthor-gpu; then
  :
elif grep -q '^overlays=' "$env_file"; then
  sed -i '/^overlays=/ s/$/ panthor-gpu/' "$env_file"
else
  printf '\noverlays=panthor-gpu\n' >>"$env_file"
fi
mkdir -p "$STAGE_DIR/boot"
cp "$env_file" "$STAGE_DIR/boot/armbianEnv.txt"

# Host umasks must not leak group-writable modes into /etc or systemd paths.
find "$STAGE_DIR" -type d -exec chmod 0755 {} +
find "$STAGE_DIR" -type f ! -perm /0111 -exec chmod 0644 {} +
find "$STAGE_DIR" -type f -perm /0111 -exec chmod 0755 {} +
chmod 0640 "$STAGE_DIR/etc/shadow"

"${SCRIPT_DIR}/inject-ext4-tree.sh" "$ROOTFS_IMAGE" "$STAGE_DIR"
debugfs -w -R 'set_inode_field /etc/shadow gid 42' "$ROOTFS_IMAGE" >/dev/null
debugfs -w -R 'set_inode_field /etc/shadow mode 0100640' "$ROOTFS_IMAGE" >/dev/null

# Replace the overlay's portable placeholder files with real enablement links.
debugfs -w -R 'rm /etc/systemd/system/sysinit.target.wants/rk1-media-identity.service' "$ROOTFS_IMAGE" >/dev/null 2>&1 || true
debugfs -w -R 'symlink /etc/systemd/system/sysinit.target.wants/rk1-media-identity.service ../rk1-media-identity.service' "$ROOTFS_IMAGE" >/dev/null
debugfs -w -R 'rm /etc/systemd/system/multi-user.target.wants/rk1-media-provision.service' "$ROOTFS_IMAGE" >/dev/null 2>&1 || true
debugfs -w -R 'symlink /etc/systemd/system/multi-user.target.wants/rk1-media-provision.service ../rk1-media-provision.service' "$ROOTFS_IMAGE" >/dev/null

for stale_path in \
  /root/.not_logged_in_yet \
  /root/.ssh/authorized_keys \
  /etc/ssh/ssh_host_ecdsa_key /etc/ssh/ssh_host_ecdsa_key.pub \
  /etc/ssh/ssh_host_ed25519_key /etc/ssh/ssh_host_ed25519_key.pub \
  /etc/ssh/ssh_host_rsa_key /etc/ssh/ssh_host_rsa_key.pub; do
  debugfs -w -R "rm ${stale_path}" "$ROOTFS_IMAGE" >/dev/null 2>&1 || true
done

ext4_stat() {
  debugfs -R "stat $1" "$ROOTFS_IMAGE" 2>/dev/null
}
ext4_exists() {
  ext4_stat "$1" | grep -q '^Inode:'
}
ext4_missing() {
  ! ext4_exists "$1"
}
assert_ext4_contains() {
  local path=$1 pattern=$2
  debugfs -R "cat $path" "$ROOTFS_IMAGE" 2>/dev/null | grep -Eq "$pattern" || {
    echo "Image assertion failed for $path: $pattern" >&2
    exit 1
  }
}

assert_ext4_contains /boot/armbianEnv.txt '^fdtfile=rockchip/rk3588-turing-rk1\.dtb$'
assert_ext4_contains /boot/armbianEnv.txt '^overlays=.*panthor-gpu'
assert_ext4_contains /etc/default/armbian-firstrun '^OPENSSHD_REGENERATE_HOST_KEYS=false$'
assert_ext4_contains /etc/ssh/sshd_config.d/90-rk1-media.conf '^AuthenticationMethods publickey$'
assert_ext4_contains /etc/shadow '^root:!:'
assert_ext4_contains /usr/local/sbin/rk1-media-provision \
  '^seed_archives=/opt/rk1-seed/debs$'
assert_ext4_contains /usr/local/sbin/rk1-media-provision \
  'Dir::Cache::archives=.*seed_archives'
assert_ext4_contains /usr/local/sbin/rk1-media-provision \
  '^systemctl --no-block start .*kodi-rk\.service$'
assert_ext4_contains /usr/local/sbin/rk1-media-provision \
  '^udevadm trigger --action=add --subsystem-match=misc'
ext4_exists /opt/rknn/2.3.2/lib/librknnrt.so || { echo 'RKNN runtime is absent' >&2; exit 1; }
ext4_exists /usr/local/sbin/rk1-media-migrate-root-to-nvme || {
  echo 'Guarded NVMe migration helper is absent' >&2
  exit 1
}
ext4_exists /opt/rk1-seed/debs/"$(basename -- "$KODI_DEB")" || { echo 'Kodi package is absent' >&2; exit 1; }
expected_deb_count=$(find "$STAGE_DIR/opt/rk1-seed/debs" -maxdepth 1 -type f -name '*.deb' | wc -l)
actual_deb_count=$(debugfs -R 'ls -p /opt/rk1-seed/debs' "$ROOTFS_IMAGE" 2>/dev/null | grep -c '/.*\.deb/' || true)
[[ "$actual_deb_count" == "$expected_deb_count" ]] || {
  echo "Offline package bundle is incomplete ($actual_deb_count/$expected_deb_count)" >&2
  exit 1
}
ext4_stat /etc | grep -Eq 'Mode:[[:space:]]+0755' || { echo '/etc mode is unsafe' >&2; exit 1; }
ext4_stat /etc/systemd/system/sysinit.target.wants/rk1-media-identity.service \
  | grep -Fq 'Fast link dest: "../rk1-media-identity.service"' || { echo 'Identity service is not enabled' >&2; exit 1; }
ext4_stat /etc/systemd/system/multi-user.target.wants/rk1-media-provision.service \
  | grep -Fq 'Fast link dest: "../rk1-media-provision.service"' || { echo 'Provisioning service is not enabled' >&2; exit 1; }
for forbidden_path in /root/.not_logged_in_yet /root/.ssh/authorized_keys \
  /etc/ssh/ssh_host_ecdsa_key /etc/ssh/ssh_host_ed25519_key /etc/ssh/ssh_host_rsa_key; do
  ext4_missing "$forbidden_path" || { echo "Forbidden image-time credential: $forbidden_path" >&2; exit 1; }
done

repair_ext4 "$ROOTFS_IMAGE"
dd if="$ROOTFS_IMAGE" of="$RAW_IMAGE" bs=512 seek="$partition_start" conv=notrunc status=none
sgdisk --verify "$RAW_IMAGE" >/dev/null

cp --reflink=auto "$RAW_IMAGE" "${OUTPUT_DIR}/${OUTPUT_BASENAME}.img"
xz --threads=0 -3 --keep --force "${OUTPUT_DIR}/${OUTPUT_BASENAME}.img"
(
  cd "$OUTPUT_DIR"
  sha256sum "${OUTPUT_BASENAME}.img.xz" >"${OUTPUT_BASENAME}.img.xz.sha256"
)

raw_sha=$(sha256sum "${OUTPUT_DIR}/${OUTPUT_BASENAME}.img" | awk '{print $1}')
compressed_sha=$(sha256sum "${OUTPUT_DIR}/${OUTPUT_BASENAME}.img.xz" | awk '{print $1}')
fingerprints_json=$(ssh-keygen -lf "$SSH_PUBLIC_KEY" | awk '{print $2}' \
  | jq -Rsc 'split("\n") | map(select(length > 0))')
jq -n \
  --arg build_id "$BUILD_ID" \
  --arg base_file "$(basename -- "$BASE_XZ")" \
  --arg base_sha256 "$actual_base_sha" \
  --arg media_deb "$(basename -- "$KODI_DEB")" \
  --arg media_deb_sha256 "$(sha256sum "$KODI_DEB" | awk '{print $1}')" \
  --arg rknn_version 2.3.2 \
  --arg rknn_commit 42aa1d426c0a9e0869b6374edba009f7208a1926 \
  --arg inputs_lock_sha256 "$(sha256sum "$PROJECT_DIR/inputs.lock.json" | awk '{print $1}')" \
  --arg debian_lock_sha256 "$(sha256sum "$PROJECT_DIR/packages/debian-packages.lock.tsv" | awk '{print $1}')" \
  --arg raw_sha256 "$raw_sha" \
  --arg compressed_sha256 "$compressed_sha" \
  --argjson image_size_mib "$IMAGE_SIZE_MIB" \
  --argjson dependency_package_count "$expected_deb_count" \
  --argjson ssh_key_fingerprints "$fingerprints_json" \
  '{
    schema_version: 1,
    build_id: $build_id,
    board: "turing-rk1",
    release: "debian-trixie-13.6",
    kernel: "6.1.115-vendor-rk35xx",
    gpu_stack: "panthor+mesa-panfrost+panvk",
    base: {file: $base_file, sha256: $base_sha256},
    media_package: {file: $media_deb, sha256: $media_deb_sha256},
    rknn: {version: $rknn_version, commit: $rknn_commit},
    locks: {inputs_sha256: $inputs_lock_sha256, debian_packages_sha256: $debian_lock_sha256},
    artifacts: {
      raw_sha256: $raw_sha256,
      compressed_sha256: $compressed_sha256,
      image_size_mib: $image_size_mib
    },
    dependency_package_count: $dependency_package_count,
    ssh_key_fingerprints: $ssh_key_fingerprints
  }' >"${OUTPUT_DIR}/${OUTPUT_BASENAME}.manifest.json"

echo "Built ${OUTPUT_DIR}/${OUTPUT_BASENAME}.img.xz"