#!/usr/bin/env bash
# Run on the Proxmox host when build_cloudinit_request.json appears in the provisioning dir.
# Downloads Raspberry Pi OS image from URL, injects cloud-init NoCloud files, saves to BACKUPS_DIR.
# Uses loop devices and mount (available on host, not in unprivileged LXC).
# Triggered by systemd path unit cm4-build-cloudinit.path.

set -e
PROV_DIR="${CM4_PROVISIONING_DIR:-/var/lib/cm4-provisioning}"
REQUEST_FILE="$PROV_DIR/build_cloudinit_request.json"
STATUS_FILE="$PROV_DIR/build_cloudinit_status.json"
CANCEL_FILE="$PROV_DIR/build_cloudinit_cancel"
[[ -f /opt/cm4-provisioning/env ]] && source /opt/cm4-provisioning/env
BACKUPS_DIR="${BACKUPS_DIR:-$PROV_DIR/backups}"
CLOUDINIT_IMAGES_DIR="${CLOUDINIT_IMAGES_DIR:-$PROV_DIR/cloudinit-images}"
DOWNLOAD_CACHE_DIR="${CM4_DOWNLOAD_CACHE_DIR:-$PROV_DIR/download-cache}"
# Write built cloud-init images to CLOUDINIT_IMAGES_DIR (separate from backups)
OUTPUT_DIR="${CLOUDINIT_IMAGES_DIR}"
GOLDEN_IMAGE="${GOLDEN_IMAGE:-$PROV_DIR/golden.img}"

write_status() {
    local phase="$1" message="$2" output_name="$3" error="$4"
    printf '{"phase":"%s","message":"%s","output_name":%s,"error":%s,"updated":%s}\n' \
        "$phase" "$message" \
        "$([ -n "$output_name" ] && echo "\"$output_name\"" || echo "null")" \
        "$([ -n "$error" ] && echo "\"${error//\"/\\\"}\"" || echo "null")" \
        "$(date +%s)" > "$STATUS_FILE"
}

check_cancel() {
    if [[ -f "$CANCEL_FILE" ]]; then
        write_status "cancelled" "Build cancelled." "" ""
        rm -f "$REQUEST_FILE" "$CANCEL_FILE"
        exit 0
    fi
}

[[ -f "$REQUEST_FILE" ]] || { echo "No request file"; exit 0; }
check_cancel

# Use temp dir on provisioning dir (not /tmp) so we have enough space for decompress (~3GB+)
mkdir -p "$PROV_DIR"
TEMP_DIR=$(mktemp -d -p "$PROV_DIR" build.XXXXXX 2>/dev/null) || TEMP_DIR="$PROV_DIR/build.$$" && mkdir -p "$TEMP_DIR"
trap 'rm -rf "$TEMP_DIR"; rm -f "$REQUEST_FILE"' EXIT

# Extract fields from JSON into temp files (handles multi-line content)
python3 - "$REQUEST_FILE" "$TEMP_DIR" << 'PY'
import json, sys
with open(sys.argv[1]) as f:
    d = json.load(f)
out = sys.argv[2]
for name, key in [("user-data", "user_data"), ("meta-data", "meta_data"), ("network-config", "network_config")]:
    with open(f"{out}/{name}", "w") as f:
        f.write(d.get(key, ""))
with open(f"{out}/variant", "w") as f:
    f.write(d.get("variant", "lite"))
with open(f"{out}/url", "w") as f:
    f.write(d.get("url", ""))
with open(f"{out}/set_golden", "w") as f:
    f.write("1" if d.get("set_as_golden_after") else "0")
name = (d.get("image_name") or "").strip()[:64]
name = "".join(c for c in name if c.isalnum() or c in "_-")
with open(f"{out}/image_name", "w") as f:
    f.write(name)
PY

URL=$(cat "$TEMP_DIR/url")
VARIANT=$(cat "$TEMP_DIR/variant")
SET_AS_GOLDEN=$(cat "$TEMP_DIR/set_golden")
IMAGE_NAME="$(cat "$TEMP_DIR/image_name" 2>/dev/null || true)"
[[ -n "$URL" ]] || { write_status "error" "" "" "Missing url in request"; exit 1; }

DATE_SUFFIX="$(date +%Y%m%d-%H%M%S)"
if [[ -n "$IMAGE_NAME" ]]; then
  OUT_NAME="${IMAGE_NAME}-${DATE_SUFFIX}.img"
else
  OUT_NAME="raspios-${VARIANT}-cloudinit-${DATE_SUFFIX}.img"
fi
OUT_PATH="$OUTPUT_DIR/$OUT_NAME"

BASENAME=$(basename "$URL")
SHA256_URL="${URL}.sha256"
CACHED="$DOWNLOAD_CACHE_DIR/$BASENAME"
XZ_FILE="$TEMP_DIR/image.img.xz"
CURL_ERR="$TEMP_DIR/curl_err.txt"
SHA256_FILE="$TEMP_DIR/expected.sha256"

# Fetch official SHA256 checksum if available (Raspberry Pi OS publishes .img.xz.sha256 next to each image)
check_cancel
write_status "downloading" "Checking for existing image and checksum…" "" ""
EXPECTED_HASH=""
if curl -sfL -o "$SHA256_FILE" "$SHA256_URL" 2>/dev/null && [[ -s "$SHA256_FILE" ]]; then
    # Format is typically "hash  filename" or just "hash"; take first 64-char hex token
    EXPECTED_HASH=$(head -1 "$SHA256_FILE" | awk '{print $1}')
    [[ ${#EXPECTED_HASH} -eq 64 && "$EXPECTED_HASH" =~ ^[0-9a-fA-F]+$ ]] || EXPECTED_HASH=""
fi

USE_CACHED=0
mkdir -p "$DOWNLOAD_CACHE_DIR"
if [[ -f "$CACHED" && -s "$CACHED" ]]; then
    if [[ -n "$EXPECTED_HASH" ]]; then
        ACTUAL_HASH=$(sha256sum -b "$CACHED" 2>/dev/null | awk '{print $1}')
        if [[ "$ACTUAL_HASH" == "$EXPECTED_HASH" ]]; then
            USE_CACHED=1
        else
            rm -f "$CACHED"
        fi
    else
        USE_CACHED=1
    fi
fi

if [[ $USE_CACHED -eq 1 ]]; then
    write_status "downloading" "Using cached image (checksum OK): $BASENAME" "" ""
    cp "$CACHED" "$XZ_FILE"
else
    write_status "downloading" "Downloading $(basename "$URL")…" "" ""
    CURL_OPTS=(-f -L -o "$XZ_FILE" --connect-timeout 30 --max-time 7200)
    [[ "${CURL_INSECURE:-}" == "1" ]] && CURL_OPTS+=(-k)
    if ! curl "${CURL_OPTS[@]}" "$URL" 2>"$CURL_ERR"; then
        err_detail=$(head -c 200 "$CURL_ERR" | tr '\n' ' ' | sed 's/"/\\"/g')
        write_status "error" "" "" "Download failed: ${err_detail:-curl exited with error}"
        exit 1
    fi
    # Verify with official checksum if we have it
    if [[ -n "$EXPECTED_HASH" ]]; then
        ACTUAL_HASH=$(sha256sum -b "$XZ_FILE" 2>/dev/null | awk '{print $1}')
        if [[ "$ACTUAL_HASH" != "$EXPECTED_HASH" ]]; then
            write_status "error" "" "" "Checksum mismatch: download may be corrupted (expected ${EXPECTED_HASH:0:16}..., got ${ACTUAL_HASH:0:16}...). Try again."
            exit 1
        fi
        write_status "downloading" "Checksum OK, caching image…" "" ""
    fi
    cp "$XZ_FILE" "$CACHED"
fi

check_cancel
write_status "decompressing" "Decompressing image…" "" ""
# Check we have a real xz file (not HTML error page)
if ! command -v xz >/dev/null 2>&1; then
    write_status "error" "" "" "Decompress failed: xz not installed. Install xz-utils (apt install xz-utils)"
    exit 1
fi
if [[ ! -s "$XZ_FILE" ]]; then
    write_status "error" "" "" "Decompress failed: downloaded file is empty"
    exit 1
fi
FILE_TYPE=$(file -b "$XZ_FILE" 2>/dev/null || true)
if [[ "$FILE_TYPE" == *"HTML"* ]] || [[ "$FILE_TYPE" == *"text"* ]] && [[ "$FILE_TYPE" != *"XZ"* ]]; then
    write_status "error" "" "" "Decompress failed: download is not an image (got: ${FILE_TYPE:0:80})"
    exit 1
fi
# Need ~3GB+ free for decompressing typical Raspios .img.xz; use same filesystem as TEMP_DIR
AVAIL_KB=$(df -P "$TEMP_DIR" 2>/dev/null | awk 'NR==2 {print $4}')
AVAIL_GB=0
[[ -n "$AVAIL_KB" ]] && AVAIL_GB=$((AVAIL_KB / 1024 / 1024))
if [[ -n "$AVAIL_KB" && "$AVAIL_KB" -lt 3145728 ]]; then
    write_status "error" "" "" "Decompress failed: not enough disk space (${AVAIL_GB}GB free on $(df -P "$TEMP_DIR" 2>/dev/null | awk 'NR==2 {print $6}'); need ~3GB). Free space or use a larger volume."
    exit 1
fi
IMG_FILE="$TEMP_DIR/image.img"
XZ_ERR="$TEMP_DIR/xz_err.txt"
# -T 1: single-threaded to reduce memory use; -d -k -f: decompress, keep source, force
if ! xz -T 1 -d -k -f "$XZ_FILE" 2>"$XZ_ERR"; then
    err_detail=$(cat "$XZ_ERR" 2>/dev/null | tr '\n' ' ' | sed 's/"/\\"/g' | head -c 400)
    [[ -z "$err_detail" ]] && err_detail="out of disk space or memory?"
    write_status "error" "" "" "Decompress failed: ${err_detail} (free space: ${AVAIL_GB}GB)"
    exit 1
fi
[[ -f "$IMG_FILE" ]] || { write_status "error" "" "" "image.img not found after decompress"; exit 1; }

check_cancel
write_status "injecting" "Mounting boot partition and injecting cloud-init…" "" ""
LOOP=$(losetup -f --show -P "$IMG_FILE")
boot_part="${LOOP}p1"
[[ -b "$boot_part" ]] || boot_part="${LOOP}p2"
[[ -b "$boot_part" ]] || { write_status "error" "" "" "Boot partition not found"; losetup -d "$LOOP"; exit 1; }

MNT="$TEMP_DIR/mnt"
mkdir -p "$MNT"
mount "$boot_part" "$MNT"
cp "$TEMP_DIR/user-data" "$MNT/user-data"
cp "$TEMP_DIR/meta-data" "$MNT/meta-data"
cp "$TEMP_DIR/network-config" "$MNT/network-config"
umount "$MNT"
losetup -d "$LOOP"

check_cancel
write_status "finalizing" "Copying image to cloud-init images…" "" ""
mkdir -p "$OUTPUT_DIR"
cp "$IMG_FILE" "$OUT_PATH"

# Compress to .img.xz to reduce size (deploy supports decompress on the fly)
write_status "finalizing" "Compressing image (xz)…" "" ""
OUT_XZ="${OUT_PATH}.xz"
# -T 0 = use all cores; fallback to -T 1 if old xz
if ! xz -T 0 -z -k -f "$OUT_PATH" 2>/dev/null; then
    xz -T 1 -z -k -f "$OUT_PATH" 2>/dev/null || true
fi
if [[ -f "$OUT_XZ" ]]; then
    rm -f "$OUT_PATH"
    OUT_NAME="${OUT_NAME}.xz"
    OUT_PATH="$OUT_XZ"
else
    write_status "finalizing" "Compression skipped (xz failed), keeping raw image." "" ""
fi

if [[ "$SET_AS_GOLDEN" == "1" ]]; then
    rm -f "$GOLDEN_IMAGE"
    ln -sf "$OUT_PATH" "$GOLDEN_IMAGE"
    write_status "done" "Built $OUT_NAME and set as golden image." "$OUT_NAME" ""
else
    write_status "done" "Built $OUT_NAME" "$OUT_NAME" ""
fi