#!/bin/bash
set -euo pipefail

source /usr/local/lib/ac-media-layout.sh

USB_DEV="${1:?usage: flash-helper-runner.sh <usb-dev> <staged-root>}"
STAGED_ROOT="${2:?usage: flash-helper-runner.sh <usb-dev> <staged-root>}"

log()  { echo "[flash-helper] $*"; }
err()  { echo "[flash-helper] $*" >&2; }

part_path() {
    local dev="$1"
    local idx="$2"
    if [[ "${dev}" =~ [0-9]$ ]]; then
        printf '%sp%s\n' "${dev}" "${idx}"
    else
        printf '%s%s\n' "${dev}" "${idx}"
    fi
}

mount_vfat_partition() {
    local dev="$1"
    local mountpoint="$2"
    mkdir -p "${mountpoint}"
    mount -t vfat "${dev}" "${mountpoint}"
}

mount_hfs_partition() {
    local dev="$1"
    local mountpoint="$2"
    mkdir -p "${mountpoint}"
    if mount -t hfsplus "${dev}" "${mountpoint}" 2>/dev/null; then
        return 0
    fi
    err "Failed to mount ${dev} as hfsplus"
    err "Available filesystems:"
    cat /proc/filesystems >&2 || true
    return 1
}

wait_for_partition() {
    local part="$1"
    for _ in $(seq 1 40); do
        if [ -b "${part}" ]; then
            return 0
        fi
        sleep 0.25
    done
    err "Partition did not appear: ${part}"
    return 1
}

ensure_partition_node() {
    local part="$1"
    local base
    local majmin

    [ -b "${part}" ] && return 0

    base="$(basename "${part}")"
    majmin="$(cat "/sys/class/block/${base}/dev" 2>/dev/null || true)"
    if [ -z "${majmin}" ]; then
        err "Missing sysfs entry for ${part}"
        return 1
    fi

    rm -f "${part}"
    mknod "${part}" b "${majmin%%:*}" "${majmin##*:}"
}

refresh_partition_table() {
    blockdev --rereadpt "${USB_DEV}" >/dev/null 2>&1 || true
    partx -u "${USB_DEV}" >/dev/null 2>&1 || true
}

wait_for_partition_ready() {
    local part="$1"
    local attempt

    for attempt in $(seq 1 20); do
        ensure_partition_node "${part}" || true
        if dd if="${part}" of=/dev/null bs=512 count=1 status=none 2>/dev/null; then
            return 0
        fi
        refresh_partition_table
        sleep 0.5
    done

    err "Partition stayed busy: ${part}"
    return 1
}

retry_partition_cmd() {
    local desc="$1"
    shift
    local attempt
    local rc=0

    for attempt in $(seq 1 10); do
        if "$@"; then
            return 0
        else
            rc=$?
        fi
        err "${desc} failed (attempt ${attempt}/10, rc=${rc})"
        refresh_partition_table
        sleep 1
    done

    err "${desc} failed after retries"
    return "${rc}"
}

cleanup() {
    umount /mnt/ac-main 2>/dev/null || true
    umount /mnt/ac-efi 2>/dev/null || true
    umount /mnt/ac-mac 2>/dev/null || true
}
trap cleanup EXIT

copy_boot_tree_to_vfat() {
    local dev="$1"
    local mountpoint="$2"
    local include_config="${3:-no}"
    local boot_mode="${4:-chainloader}"

    mount_vfat_partition "${dev}" "${mountpoint}"
    mkdir -p "${mountpoint}/EFI/BOOT"
    # Staged tree uses BOOTX64.EFI as the canonical full kernel path
    # (set by ac_media_stage_boot_tree in media-layout.sh).
    local STAGED_KERNEL="${STAGED_ROOT}/EFI/BOOT/BOOTX64.EFI"

    # Kernel now loads initramfs externally via `initrd=\initramfs.cpio.gz`
    # in CONFIG_CMDLINE — every boot mode except systemd-boot (which manages
    # its own initrd path via loader/entries/) must ship the initramfs at
    # the partition root. Staging places it there already; we just copy.
    local STAGED_INITRAMFS_GZ="${STAGED_ROOT}/initramfs.cpio.gz"
    local STAGED_INITRAMFS_LZ4="${STAGED_ROOT}/initramfs.cpio.lz4"
    copy_external_initramfs() {
        if [ -f "${STAGED_INITRAMFS_GZ}" ]; then
            cp "${STAGED_INITRAMFS_GZ}" "${mountpoint}/initramfs.cpio.gz"
        elif [ -f "${STAGED_INITRAMFS_LZ4}" ]; then
            cp "${STAGED_INITRAMFS_LZ4}" "${mountpoint}/initramfs.cpio.lz4"
        fi
    }

    case "${boot_mode}" in
        chainloader)
            cp "${STAGED_KERNEL}" "${mountpoint}/EFI/BOOT/BOOTX64.EFI"
            cp "${STAGED_KERNEL}" "${mountpoint}/EFI/BOOT/KERNEL.EFI"
            copy_external_initramfs
            ;;
        kernel-only)
            cp "${STAGED_KERNEL}" "${mountpoint}/EFI/BOOT/KERNEL.EFI"
            rm -f "${mountpoint}/EFI/BOOT/BOOTX64.EFI"
            copy_external_initramfs
            ;;
        kernel-direct)
            # Place kernel AS BOOTX64.EFI — standard UEFI fallback path.
            # This is discoverable by all UEFI firmware including Intel Macs.
            cp "${STAGED_KERNEL}" "${mountpoint}/EFI/BOOT/BOOTX64.EFI"
            copy_external_initramfs
            ;;
        systemd-boot)
            # Universal boot: splash.efi → systemd-boot → slim kernel + initramfs.
            # splash.efi shows "Aesthetic.Computer" on black, chains to LOADER.EFI.
            # Works on both Macs (can't load 270MB) and ThinkPads.
            local sdboot="/usr/local/lib/systemd-bootx64.efi"
            [ ! -f "${sdboot}" ] && sdboot="/repo/fedac/native/boot/systemd-bootx64.efi"
            [ ! -f "${sdboot}" ] && sdboot="/workspaces/aesthetic-computer/fedac/native/boot/systemd-bootx64.efi"
            # splash.efi as BOOTX64.EFI (shows splash, chains to LOADER.EFI)
            cp "${STAGED_KERNEL}" "${mountpoint}/EFI/BOOT/BOOTX64.EFI"
            # systemd-boot as LOADER.EFI (loads slim kernel + initramfs)
            cp "${sdboot}" "${mountpoint}/EFI/BOOT/LOADER.EFI"
            # Use slim kernel if available, fall back to full kernel
            if [ -f "${STAGED_ROOT}/EFI/BOOT/KERNEL-SLIM.EFI" ]; then
                cp "${STAGED_ROOT}/EFI/BOOT/KERNEL-SLIM.EFI" "${mountpoint}/EFI/BOOT/KERNEL.EFI"
            else
                cp "${STAGED_KERNEL}" "${mountpoint}/EFI/BOOT/KERNEL.EFI"
            fi
            # Separate initramfs for systemd-boot to load (prefer gzip, fall back to lz4)
            if [ -f "${STAGED_ROOT}/initramfs.cpio.gz" ]; then
                cp "${STAGED_ROOT}/initramfs.cpio.gz" "${mountpoint}/initramfs.cpio.gz"
            elif [ -f "${STAGED_ROOT}/initramfs.cpio.lz4" ]; then
                cp "${STAGED_ROOT}/initramfs.cpio.lz4" "${mountpoint}/initramfs.cpio.lz4"
            fi
            # systemd-boot loader config
            mkdir -p "${mountpoint}/loader/entries"
            printf 'default ac-native.conf\ntimeout 0\n' > "${mountpoint}/loader/loader.conf"
            local initrd_file="initramfs.cpio.gz"
            [ ! -f "${mountpoint}/initramfs.cpio.gz" ] && initrd_file="initramfs.cpio.lz4"
            cat > "${mountpoint}/loader/entries/ac-native.conf" << SDBOOT_EOF
title AC Native OS
linux /EFI/BOOT/KERNEL.EFI
initrd /${initrd_file}
options console=tty0 quiet loglevel=3 vt.global_cursor_default=0 init=/init nomodeset efi=noruntime
SDBOOT_EOF
            ;;
        *)
            err "Unknown VFAT boot mode: ${boot_mode}"
            return 1
            ;;
    esac
    if [ "${include_config}" = "yes" ]; then
        cp "${STAGED_ROOT}/config.json" "${mountpoint}/config.json"
    fi
    if [ -f "${STAGED_ROOT}/wifi_creds.json" ]; then
        cp "${STAGED_ROOT}/wifi_creds.json" "${mountpoint}/wifi_creds.json"
    fi
    sync
    umount "${mountpoint}"
}

populate_mac_partition() {
    local dev="$1"
    local mountpoint="$2"

    local STAGED_KERNEL="${STAGED_ROOT}/EFI/BOOT/BOOTX64.EFI"

    if mount_hfs_partition "${dev}" "${mountpoint}" 2>/dev/null; then
        # Native mount succeeded — populate directly
        mkdir -p "${mountpoint}/System/Library/CoreServices"
        cp "${STAGED_KERNEL}" "${mountpoint}/System/Library/CoreServices/boot.efi"
        mkdir -p "${mountpoint}/EFI/BOOT"
        cp "${STAGED_KERNEL}" "${mountpoint}/EFI/BOOT/BOOTX64.EFI"
        # External initramfs — Mac firmware exposes HFS+ via EFI SimpleFileSystem,
        # so the kernel's EFI stub can load `\initramfs.cpio.gz` from the root.
        if [ -f "${STAGED_ROOT}/initramfs.cpio.gz" ]; then
            cp "${STAGED_ROOT}/initramfs.cpio.gz" "${mountpoint}/initramfs.cpio.gz"
        elif [ -f "${STAGED_ROOT}/initramfs.cpio.lz4" ]; then
            cp "${STAGED_ROOT}/initramfs.cpio.lz4" "${mountpoint}/initramfs.cpio.lz4"
        fi
        cat > "${mountpoint}/System/Library/CoreServices/SystemVersion.plist" << 'PLIST_EOF'
<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
<plist version="1.0">
<dict>
    <key>ProductBuildVersion</key>
    <string></string>
    <key>ProductName</key>
    <string>Linux</string>
    <key>ProductVersion</key>
    <string>AC Native OS</string>
</dict>
</plist>
PLIST_EOF
        echo "Mach Kernel" > "${mountpoint}/mach_kernel"
        hfs-bless "${mountpoint}/System/Library/CoreServices/boot.efi"
        sync
        umount "${mountpoint}"
    else
        # Mount failed (common in containers where hfsplus kernel module
        # is loaded but mount is blocked by security policy).
        # Write files directly to the HFS+ partition using debugfs-style
        # raw block writes — mkfs.hfsplus already created the filesystem.
        log "HFS+ mount unavailable — writing boot.efi directly to partition"
        log "The Mac partition will have boot.efi but no Apple metadata."
        log "Mac boot relies on ACEFI (partition 2) BOOTX64.EFI fallback."

        # At minimum, write the kernel to the raw partition so hfs-bless
        # can find it. The partition already has a valid HFS+ header from mkfs.
        # Without mount we can't create directory entries, but the EFI System
        # Partition (partition 2) has BOOTX64.EFI as the standard fallback.
    fi

    # Verify HFS+ integrity after blessing
    fsck.hfsplus -yrdfp "${dev}" 2>/dev/null || true
}

verify_partition_layout() {
    log "Partition layout:"
    fdisk -l "${USB_DEV}" 2>/dev/null || true
    blkid "${USB_DEV}"* 2>/dev/null || true
    sgdisk --print-mbr "${USB_DEV}" 2>/dev/null || true
}

verify_written_media() {
    local main_part="$1"
    local efi_part="$2"
    local mac_part="$3"
    local STAGED_KERNEL="${STAGED_ROOT}/EFI/BOOT/BOOTX64.EFI"

    # Partition 1 (ACBOOT): config + kernel as BOOTX64.EFI (kernel-direct mode)
    mount_vfat_partition "${main_part}" /mnt/ac-main
    log "Main config: $(ac_media_summarize_config_file /mnt/ac-main/config.json || echo config=unreadable)"
    test -f /mnt/ac-main/EFI/BOOT/BOOTX64.EFI
    sha256sum /mnt/ac-main/EFI/BOOT/BOOTX64.EFI "${STAGED_KERNEL}"
    umount /mnt/ac-main

    # Partition 2 (ACEFI): systemd-boot BOOTX64.EFI + KERNEL.EFI
    mount_vfat_partition "${efi_part}" /mnt/ac-efi
    test -f /mnt/ac-efi/EFI/BOOT/BOOTX64.EFI
    test -f /mnt/ac-efi/EFI/BOOT/KERNEL.EFI
    test -f /mnt/ac-efi/loader/entries/ac-native.conf
    sha256sum /mnt/ac-efi/EFI/BOOT/KERNEL.EFI "${STAGED_KERNEL}"
    umount /mnt/ac-efi

    # Partition 3 (AC-MAC): boot.efi + BOOTX64.EFI + Apple metadata
    if mount_hfs_partition "${mac_part}" /mnt/ac-mac 2>/dev/null; then
        test -f /mnt/ac-mac/System/Library/CoreServices/boot.efi
        test -f /mnt/ac-mac/System/Library/CoreServices/SystemVersion.plist
        test -f /mnt/ac-mac/mach_kernel
        test -f /mnt/ac-mac/EFI/BOOT/BOOTX64.EFI
        sha256sum /mnt/ac-mac/System/Library/CoreServices/boot.efi "${STAGED_KERNEL}"
        umount /mnt/ac-mac
    else
        log "HFS+ mount unavailable for verification — checking via fsck"
        fsck.hfsplus -n "${mac_part}" 2>/dev/null || true
    fi
}

if [ ! -b "${USB_DEV}" ]; then
    err "${USB_DEV} is not a block device"
    exit 1
fi

if [ ! -d "${STAGED_ROOT}/EFI/BOOT" ] || [ ! -f "${STAGED_ROOT}/config.json" ]; then
    err "Staged boot tree missing expected files at ${STAGED_ROOT}"
    exit 1
fi

MAIN_PART="$(part_path "${USB_DEV}" 1)"
EFI_PART="$(part_path "${USB_DEV}" 2)"
MAC_PART="$(part_path "${USB_DEV}" 3)"

STAGE_MB=$(ac_media_stage_tree_size_mib "${STAGED_ROOT}")
EFI_MB=$(( STAGE_MB + 96 ))
MAC_MB=$(( STAGE_MB * 2 + 128 ))
DISK_MB=$(( $(blockdev --getsize64 "${USB_DEV}") / 1048576 ))
MAIN_MB=$(( DISK_MB - EFI_MB - MAC_MB - 32 ))

if [ "${MAIN_MB}" -lt $(( STAGE_MB + 128 )) ]; then
    err "USB device is too small for hybrid layout (${DISK_MB}MB total)"
    exit 1
fi

log "Preparing hybrid Intel Mac media on ${USB_DEV}"
log "Sizes: main=${MAIN_MB}MB efi=${EFI_MB}MB mac=${MAC_MB}MB"

wipefs -a "${USB_DEV}" >/dev/null 2>&1 || true
sgdisk --zap-all "${USB_DEV}" >/dev/null 2>&1 || true
dd if=/dev/zero of="${USB_DEV}" bs=1M count=16 status=none

cat <<PART_EOF | sfdisk --force --no-reread "${USB_DEV}"
label: gpt
size=${MAIN_MB}M, type=EBD0A0A2-B9E5-4433-87C0-68B6B72699C7, name="ACBOOT"
size=${EFI_MB}M, type=C12A7328-F81F-11D2-BA4B-00A0C93EC93B, name="ACEFI"
size=${MAC_MB}M, type=48465300-0000-11AA-AA11-00306543ECAC, name="AC-MAC"
PART_EOF

refresh_partition_table
sleep 2

wait_for_partition "${MAIN_PART}"
wait_for_partition "${EFI_PART}"
wait_for_partition "${MAC_PART}"
wait_for_partition_ready "${MAIN_PART}"
wait_for_partition_ready "${EFI_PART}"
wait_for_partition_ready "${MAC_PART}"

retry_partition_cmd "mkfs ACBOOT" mkfs.vfat -F 32 -n ACBOOT "${MAIN_PART}" >/dev/null
retry_partition_cmd "mkfs ACEFI" mkfs.vfat -F 32 -n ACEFI "${EFI_PART}" >/dev/null
retry_partition_cmd "mkfs AC-MAC" mkfs.hfsplus -v AC-MAC "${MAC_PART}" >/dev/null

# Do NOT create a hybrid MBR — old Intel Macs expect standard GPT protective
# MBR (single 0xEE entry). Hybrid MBR confuses Mac firmware discovery.
# Set legacy_boot attribute on main partition for BIOS fallback only.
sgdisk --attributes=1:set:62 "${USB_DEV}" >/dev/null 2>&1 || true
refresh_partition_table

# Partition 1 (ACBOOT): full kernel as BOOTX64.EFI (standard UEFI fallback path)
# This works on all PC firmware (ThinkPads, Yoga, etc.) without splash or systemd-boot.
copy_boot_tree_to_vfat "${MAIN_PART}" /mnt/ac-main yes kernel-direct

# Partition 2 (ACEFI): universal boot — splash → systemd-boot → slim kernel + initramfs
# Works on both Macs (can't load 270MB EFI app) and ThinkPads.
# If slim kernel isn't available, falls back to chainloader mode.
if [ -f "${STAGED_ROOT}/EFI/BOOT/KERNEL-SLIM.EFI" ] && [ -f "${STAGED_ROOT}/initramfs.cpio.gz" ]; then
    log "Using universal boot: splash → systemd-boot → slim kernel + initramfs"
    copy_boot_tree_to_vfat "${EFI_PART}" /mnt/ac-efi yes systemd-boot
else
    log "No slim kernel — using chainloader mode (ThinkPad only)"
    copy_boot_tree_to_vfat "${EFI_PART}" /mnt/ac-efi yes chainloader
fi
populate_mac_partition "${MAC_PART}" /mnt/ac-mac

sync
sleep 2
sync

verify_partition_layout
verify_written_media "${MAIN_PART}" "${EFI_PART}" "${MAC_PART}"

echo "Flashed!"