#!/bin/bash

# ParticleOS ISO Builder with Podman - HARDENED ULTRA SAFE VERSION
# Builds a bootable ISO using Podman containers for complete isolation
# This approach cannot break your host system
# Includes all hardening features from the analysis

set -euo pipefail

# Enable logging with timestamped directory
BUILD_TIMESTAMP=$(date +"%y%m%d%H%M")
LOG_DIR="logs/${BUILD_TIMESTAMP}"
mkdir -p "$LOG_DIR"
LOG_FILE="$LOG_DIR/build-iso-podman.log"
exec > >(tee -a "$LOG_FILE") 2>&1

echo "$(date): Starting ParticleOS ISO build with logging enabled"
echo "$(date): Log directory: $LOG_DIR"

# Colors
GREEN='\033[0;32m'
BLUE='\033[0;34m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
NC='\033[0m'

print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
    # Exit immediately on error. If this is an error from the container,
    # the container's error message should have already propagated.
    exit 1 
}

print_header() {
    echo ""
    echo -e "${BLUE}================================${NC}"
    echo -e "${BLUE}$1${NC}"
    echo -e "${BLUE}================================${NC}"
}

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_NAME="particleos"
VERSION="1.0.0"
BUILD_DIR="$SCRIPT_DIR/build"
OUTPUT_DIR="$SCRIPT_DIR/output"
CONTAINER_NAME="particleos-builder"
IMAGE_NAME="particleos-builder:latest"

# Safety check: Ensure BUILD_DIR is within SCRIPT_DIR
if [[ ! "$BUILD_DIR" =~ ^"$SCRIPT_DIR" ]]; then
    print_error "BUILD_DIR ($BUILD_DIR) is not within SCRIPT_DIR ($SCRIPT_DIR). Aborting for safety."
    exit 1
fi

# Determine Podman command (sudo or not)
PODMAN_CMD=""
check_podman() {
    print_header "Phase 1: Check Podman Availability"
    
    if ! command -v podman &> /dev/null; then
        print_error "Podman is not installed. Please install Podman first: sudo apt update && sudo apt install -y podman"
    fi
    
    # Try running podman without sudo first
    if podman info &> /dev/null; then
        PODMAN_CMD="podman"
        print_success "Podman available (rootless or with user access)."
    else
        # If rootless fails, try with sudo
        if sudo podman info &> /dev/null; then
            PODMAN_CMD="sudo podman"
            print_warning "Podman requires sudo. This is normal for some configurations. Using '$PODMAN_CMD'."
        else
            print_error "Podman cannot be run even with sudo. Check your Podman installation and user permissions."
        fi
    fi
}

# Check prerequisites
check_prerequisites() {
    print_header "Phase 2: Check Prerequisites"
    
    # Check disk space (need at least 15GB free for container + ISO)
    local available_space=$(df "$SCRIPT_DIR" | awk 'NR==2 {print $4}')
    local required_space=$((15 * 1024 * 1024))   # 15GB in KB
    
    if [ "$available_space" -lt "$required_space" ]; then
        print_error "Insufficient disk space. Need at least 15GB free, have $(($available_space / 1024 / 1024))GB"
    fi
    
    # Check network connectivity
    if ! ping -c 1 archive.ubuntu.com &>/dev/null; then
        print_error "Cannot reach Ubuntu archives. Check your internet connection."
    fi
    
    # Using --silent --head for curl to check HTTP status code more reliably
    if ! curl --silent --head "https://git.raines.xyz/robojerk/apt-ostree/raw/branch/main/apt-ostree_0.1.0-1_amd64.deb" | grep "HTTP/[12] [23].." > /dev/null; then
        print_error "Cannot reach apt-ostree repository. Check your internet connection."
    fi
    
    # Check Forgejo repository accessibility (CRITICAL for bootc and ostree)
    print_status "Checking Forgejo repository accessibility..."
    if ! curl --silent "https://git.raines.xyz/api/packages/robojerk/debian/repository.key" | grep -q "BEGIN PGP PUBLIC KEY"; then
        print_error "CRITICAL: Cannot reach Forgejo repository. bootc and ostree packages are required!"
        print_error "Check if https://git.raines.xyz is accessible."
    fi
    
    # Check if Forgejo repository has packages
    if ! curl --silent "https://git.raines.xyz/api/packages/robojerk/debian" | grep -q "ostree\|bootc"; then
        print_warning "Forgejo repository may not contain ostree or bootc packages."
        print_warning "This could cause the build to fail. Proceeding anyway..."
    fi
    
    # Check apt-cacher-ng availability (optional performance enhancement)
    print_status "Checking apt-cacher-ng availability..."
    if curl --silent --head "http://192.168.1.79:3142/" | grep -q "HTTP/[12] [23].."; then
        print_success "apt-cacher-ng detected at http://192.168.1.79:3142/"
        export APT_CACHER_AVAILABLE=true
        export APT_CACHER_URL="http://192.168.1.79:3142/"
    else
        print_warning "apt-cacher-ng not available, using direct repository access"
        export APT_CACHER_AVAILABLE=false
        export APT_CACHER_URL=""
    fi
    
    print_success "All prerequisites satisfied"
}

# Clean build environment
clean_build() {
    print_header "Phase 3: Clean Build Environment"
    
    # Remove any existing container
    if $PODMAN_CMD container exists "$CONTAINER_NAME" 2>/dev/null; then
        print_status "Removing existing container: $CONTAINER_NAME"
        $PODMAN_CMD container rm -f "$CONTAINER_NAME" 2>/dev/null || true
    fi
    
    # Remove any existing image
    if $PODMAN_CMD image exists "$IMAGE_NAME" 2>/dev/null; then
        print_status "Removing existing image: $IMAGE_NAME"
        $PODMAN_CMD image rm -f "$IMAGE_NAME" 2>/dev/null || true
    fi
    
    # Clean build directories
    if [ -d "$BUILD_DIR" ]; then
        print_status "Removing previous build directory: $BUILD_DIR..."
        # Add host-side safety check for rm -rf
        if [[ "$BUILD_DIR" == "/" ]] || [[ "$BUILD_DIR" == "/home" ]] || [[ "$BUILD_DIR" == "/opt" ]] || \
           [[ "$BUILD_DIR" == "/usr" ]] || [[ "$BUILD_DIR" == "/var" ]] || [[ "$BUILD_DIR" == "/etc" ]]; then
            print_error "BUILD_DIR ($BUILD_DIR) is a critical system directory. Aborting for safety."
        fi
        rm -rf "$BUILD_DIR" || print_error "Failed to remove previous build directory."
    fi
    
    mkdir -p "$BUILD_DIR" "$OUTPUT_DIR" || print_error "Failed to create build directories."
    print_success "Build environment cleaned"
}

# Create Dockerfile for the build environment
create_dockerfile() {
    print_header "Phase 4: Create Build Environment"
    
    cat > "$BUILD_DIR/Dockerfile" << 'EOF'
FROM ubuntu:noble

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=UTC

# Update and install build tools
# Added software-properties-common for add-apt-repository
# Added live-build for more complete live ISO tools if needed, though casper/live-boot should be sufficient
# Added apt-cacher-ng for optional caching support
# Configure apt-cacher-ng if APT_CACHER_URL is provided
ARG APT_CACHER_URL
RUN if [ -n "$APT_CACHER_URL" ]; then \
        echo "Acquire::http::Proxy \"$APT_CACHER_URL\";" > /etc/apt/apt.conf.d/99proxy && \
        echo "Acquire::https::Proxy \"$APT_CACHER_URL\";" >> /etc/apt/apt.conf.d/99proxy && \
        echo "Using apt-cacher-ng: $APT_CACHER_URL"; \
    else \
        echo "Using direct repository access (no cache)"; \
    fi
RUN apt update && apt install -y \
    mmdebstrap \
    squashfs-tools \
    xorriso \
    grub-pc-bin \
    grub-efi-amd64-bin \
    isolinux \
    syslinux-common \
    syslinux-utils \
    curl \
    wget \
    ca-certificates \
    gnupg \
    software-properties-common \
    live-build \
    apt-cacher-ng \
    && rm -rf /var/lib/apt/lists/*

# Set up working directory
WORKDIR /build

# Copy build script
COPY build-in-container.sh /build/
RUN chmod +x /build/build-in-container.sh

# Default command
CMD ["/build/build-in-container.sh"]
EOF
    
    print_success "Dockerfile created"
}

# Create the build script that runs inside the container
create_container_build_script() {
    cat > "$BUILD_DIR/build-in-container.sh" << 'EOF'
#!/bin/bash

set -euo pipefail

# Colors
GREEN='\033[0;32m'
BLUE='\033[0;34m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
NC='\033[0m'

print_status() {
    echo -e "${BLUE}[CONTAINER]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[CONTAINER SUCCESS]${NC} $1"
}

print_error() {
    echo -e "${RED}[CONTAINER ERROR]${NC} $1"
    exit 1 # Exit immediately on error inside the container
}

print_warning() {
    echo -e "${YELLOW}[CONTAINER WARNING]${NC} $1"
}

print_header() {
    echo ""
    echo -e "${BLUE}================================${NC}"
    echo -e "${BLUE}$1${NC}"
    echo -e "${BLUE}================================${NC}"
}

# Configuration
PROJECT_NAME="particleos"
VERSION="1.0.0"
CHROOT_DIR="/tmp/chroot"
ISO_DIR="/tmp/iso"
OUTPUT_DIR="/output" # Mapped from host's OUTPUT_DIR

# Track mounted filesystems for cleanup inside container's chroot
MOUNTED_FILESYSTEMS_IN_CHROOT=()

# Cleanup function to unmount filesystems inside the container's chroot
cleanup_chroot_mounts() {
    print_status "Cleaning up mounted filesystems in chroot..."
    # Iterate in reverse for safer unmounting
    for (( i=${#MOUNTED_FILESYSTEMS_IN_CHROOT[@]}-1; i>=0; i-- )); do
        mount_point="${MOUNTED_FILESYSTEMS_IN_CHROOT[i]}"
        if mountpoint -q "$mount_point" 2>/dev/null; then
            print_status "Unmounting $mount_point"
            umount "$mount_point" 2>/dev/null || print_error "Failed to unmount $mount_point"
        fi
    done
    MOUNTED_FILESYSTEMS_IN_CHROOT=()
}

# Signal trap for internal script to ensure cleanup on exit/error
trap cleanup_chroot_mounts EXIT INT TERM

print_header "Starting Container Build Process"

# Create directories
mkdir -p "$CHROOT_DIR" "$ISO_DIR" "$OUTPUT_DIR" || print_error "Failed to create necessary directories in container."

# Ensure chroot directory is clean
rm -rf "$CHROOT_DIR"/* || print_error "Failed to clean chroot directory."

print_header "Phase 1: Create Base System"

# Create base system with mmdebstrap
print_status "Creating base Ubuntu system..."
mmdebstrap \
    --architectures=amd64 \
    --variant=apt \
    --include=systemd,systemd-sysv,dbus,curl,ca-certificates,gnupg,gpgv,locales,resolvconf \
    --mode=unshare \
    noble \
    "$CHROOT_DIR" \
    http://archive.ubuntu.com/ubuntu/ || print_error "mmdebstrap failed."

print_success "Base system created"

print_header "Phase 2: Configure Base System"

# Mount necessary filesystems
# Using 'mount' directly inside the container, it needs root or privileged container mode
mount --bind /dev "$CHROOT_DIR/dev" || print_error "Failed to bind mount /dev."
MOUNTED_FILESYSTEMS_IN_CHROOT+=("$CHROOT_DIR/dev")

mount --bind /run "$CHROOT_DIR/run" || print_error "Failed to bind mount /run."
MOUNTED_FILESYSTEMS_IN_CHROOT+=("$CHROOT_DIR/run")

mount -t proc none "$CHROOT_DIR/proc" || print_error "Failed to mount /proc."
MOUNTED_FILESYSTEMS_IN_CHROOT+=("$CHROOT_DIR/proc")

mount -t sysfs none "$CHROOT_DIR/sys" || print_error "Failed to mount /sys."
MOUNTED_FILESYSTEMS_IN_CHROOT+=("$CHROOT_DIR/sys")

# Fix device permissions for GPG verification
print_status "Fixing device permissions for GPG verification..."
chroot "$CHROOT_DIR" bash -c '
# Remove existing device nodes first to ensure clean creation
rm -f /dev/null /dev/zero /dev/random /dev/urandom 2>/dev/null || true

# Create device nodes with correct permissions
mknod -m 666 /dev/null c 1 3 2>/dev/null || true
mknod -m 666 /dev/zero c 1 5 2>/dev/null || true
mknod -m 644 /dev/random c 1 8 2>/dev/null || true
mknod -m 644 /dev/urandom c 1 9 2>/dev/null || true

# Ensure correct ownership (root:root)
chown root:root /dev/null /dev/zero /dev/random /dev/urandom 2>/dev/null || true

# Re-set permissions explicitly
chmod 666 /dev/null /dev/zero 2>/dev/null || true
chmod 644 /dev/random /dev/urandom 2>/dev/null || true

# Verify the device nodes exist and have correct permissions
ls -la /dev/null /dev/zero /dev/random /dev/urandom 2>/dev/null || true
' || print_error "Failed to create device nodes."

# Configure package sources
# Corrected chroot bash -c syntax
print_status "Configuring APT sources..."
chroot "$CHROOT_DIR" bash -c '
echo "deb http://archive.ubuntu.com/ubuntu noble main restricted universe multiverse" > /etc/apt/sources.list && \
echo "deb http://archive.ubuntu.com/ubuntu noble-updates main restricted universe multiverse" >> /etc/apt/sources.list && \
echo "deb http://security.ubuntu.com/ubuntu noble-security main restricted universe multiverse" >> /etc/apt/sources.list
' || print_error "Failed to configure APT sources."

# Configure apt-cacher-ng if available (optional performance enhancement)
if [ -n "${APT_CACHER_URL:-}" ]; then
    print_status "Configuring APT to use cache: $APT_CACHER_URL"
    chroot "$CHROOT_DIR" bash -c '
        echo "Acquire::http::Proxy \"$APT_CACHER_URL\";" > /etc/apt/apt.conf.d/99proxy
        echo "Acquire::https::Proxy \"$APT_CACHER_URL\";" >> /etc/apt/apt.conf.d/99proxy
    ' || print_error "Failed to configure APT cache proxy."
    print_success "APT cache configuration applied"
else
    print_status "Using direct repository access (no cache)"
fi

# Note: Mozilla PPA setup moved to host-side package download approach
# This avoids GPG verification issues inside the chroot
print_status "Mozilla PPA will be handled via host-side package download"

# Note: Forgejo repository is not accessible, so we'll skip ostree/bootc for now
print_status "Forgejo repository not accessible, skipping ostree/bootc installation"
print_status "These packages will need to be installed manually after the system is built"
print_status "apt-ostree will still be installed from direct download"

# CRITICAL FIX 3: Install gpgv first (requires --allow-unauthenticated for initial setup)
# This resolves the "gpgv, gpgv2 or gpgv1 required for verification" error.
print_status "Installing gpgv inside the chroot for APT signature verification..."
# Use --allow-unauthenticated for the initial setup since gpgv is needed for verification
chroot "$CHROOT_DIR" bash -c 'DEBIAN_FRONTEND=noninteractive apt update --allow-unauthenticated' || print_error "Failed to update APT lists for gpgv installation (pre-gpgv install)."
chroot "$CHROOT_DIR" bash -c 'DEBIAN_FRONTEND=noninteractive apt install -y --allow-unauthenticated gpgv' || print_error "Failed to explicitly install gpgv. This is required for secure package verification."
print_success "gpgv successfully installed in chroot."

# Now update package lists securely (gpgv is now available)
print_status "Updating package lists securely..."
chroot "$CHROOT_DIR" apt update || print_error "Failed to update package lists after installing gpgv. Check GPG keys and repository accessibility."

# Pre-configure problematic packages to prevent installation issues
print_status "Pre-configuring problematic packages to prevent installation failures..."
chroot "$CHROOT_DIR" bash -c '
    # Set APT preferences to prevent dictionaries-common installation BEFORE it gets installed
    echo "Package: dictionaries-common" > /etc/apt/preferences.d/99-hold-dictionaries-common.pref
    echo "Pin: release *" >> /etc/apt/preferences.d/99-hold-dictionaries-common.pref
    echo "Pin-Priority: -1" >> /etc/apt/preferences.d/99-hold-dictionaries-common.pref
    
    # Also hold aspell packages that might cause issues
    echo "Package: aspell*" >> /etc/apt/preferences.d/99-hold-dictionaries-common.pref
    echo "Pin: release *" >> /etc/apt/preferences.d/99-hold-dictionaries-common.pref
    echo "Pin-Priority: -1" >> /etc/apt/preferences.d/99-hold-dictionaries-common.pref
' || print_error "Failed to set APT preferences for problematic packages."

# Install bootc dependencies first (podman, skopeo, etc.)
print_status "Installing bootc dependencies first..."
chroot "$CHROOT_DIR" bash -c 'DEBIAN_FRONTEND=noninteractive apt install -y \
    podman \
    skopeo \
    buildah \
    conmon \
    crun \
    netavark \
    aardvark-dns \
    slirp4netns \
    fuse-overlayfs \
    uidmap \
    libsubid4 \
    libslirp0 \
    passt \
    thin-provisioning-tools \
    lvm2 \
    mdadm \
    dmeventd \
    libaio1t64 \
    libdevmapper-event1.02.1 \
    liblvm2cmd2.03 \
    kpartx \
    pigz \
    dracut \
    dracut-core' || print_error "Failed to install bootc dependencies."

# Install slimmed down desktop environment (no LibreOffice, games, etc.)
print_status "Installing slimmed down desktop environment..."
chroot "$CHROOT_DIR" bash -c 'DEBIAN_FRONTEND=noninteractive apt install -y \
    plasma-desktop \
    plasma-workspace \
    sddm \
    kwin-x11 \
    flatpak \
    network-manager \
    plasma-nm \
    openssh-server \
    curl \
    wget \
    vim \
    nano \
    htop \
    neofetch \
    tree \
    pulseaudio \
    pulseaudio-utils \
    fonts-ubuntu \
    fonts-noto \
    build-essential \
    git \
    live-boot \
    live-config \
    casper \
    systemd-sysv \
    dbus \
    locales \
    resolvconf \
    linux-image-generic \
    linux-headers-generic' || print_error "Failed to install desktop and essential packages."

# Aggressively handle problematic 'dictionaries-common' package
print_status "Attempting to fix or remove problematic 'dictionaries-common' package..."

# First, try to remove it cleanly if it's already installed and causing issues.
# Use --yes (or -y) to auto-confirm. Use --force-depends if needed.
# Capture output to see if it actually tries to remove.
if chroot "$CHROOT_DIR" dpkg -s dictionaries-common >/dev/null 2>&1; then
    print_status "'dictionaries-common' detected. Attempting to remove..."
    chroot "$CHROOT_DIR" bash -c '
        DEBIAN_FRONTEND=noninteractive apt remove -y --allow-remove-essential dictionaries-common || \
        DEBIAN_FRONTEND=noninteractive dpkg --remove --force-depends dictionaries-common || true
        apt autoremove -y
    ' || print_warning "Failed to cleanly remove dictionaries-common, attempting to hold."
fi

# Now, regardless of previous removal attempt, ensure it's held to prevent re-installation or re-configuration
# during subsequent apt operations (like apt install -f or update).
print_status "Setting APT preferences to prevent 'dictionaries-common' installation/configuration..."
chroot "$CHROOT_DIR" bash -c '
    echo "Package: dictionaries-common" > /etc/apt/preferences.d/99-hold-dictionaries-common.pref
    echo "Pin: release *" >> /etc/apt/preferences.d/99-hold-dictionaries-common.pref
    echo "Pin-Priority: -1" >> /etc/apt/preferences.d/99-hold-dictionaries-common.pref
' || print_error "Failed to set APT preferences for dictionaries-common hold."

print_status "Running apt install -f to fix any broken dependencies after package removals/holds..."
chroot "$CHROOT_DIR" apt install -f -y || print_warning "apt install -f reported issues, but continuing."

# Final check: if it's still present and broken, try to purge again
if chroot "$CHROOT_DIR" dpkg -l | grep -q "^.i.* dictionaries-common"; then
    print_warning "dictionaries-common is still present and potentially broken. Attempting final purge."
    chroot "$CHROOT_DIR" bash -c '
        DEBIAN_FRONTEND=noninteractive apt purge -y --allow-remove-essential dictionaries-common || \
        DEBIAN_FRONTEND=noninteractive dpkg --purge --force-depends dictionaries-common || true
        apt autoremove -y
    ' || print_error "CRITICAL: Could not purge dictionaries-common. Manual intervention may be required."
fi

print_success "'dictionaries-common' issue addressed (removed or held)."

# Copy pre-downloaded packages from host system
print_status "Copying pre-downloaded packages from host system..."

# Define BUILD_DIR for container script
BUILD_DIR="/build"
mkdir -p "$BUILD_DIR/packages" || print_error "Failed to create packages directory."

# Copy packages from host pkgs directory if they exist
if [ -d "/host/pkgs" ]; then
    print_status "Copying packages from /host/pkgs directory..."
    cp /host/pkgs/*.deb "$BUILD_DIR/packages/" 2>/dev/null || print_warning "Failed to copy some packages from /host/pkgs"
    print_status "Copied packages:"
    ls -la "$BUILD_DIR/packages/" || true
else
    print_warning "No /host/pkgs directory found, skipping pre-downloaded packages"
fi

# Install downloaded ostree and bootc packages (if available)
if [ -d "$BUILD_DIR/packages" ] && [ "$(ls -A "$BUILD_DIR/packages" 2>/dev/null)" ]; then
    print_status "Installing ostree, bootc, and apt-ostree packages using apt for proper dependency resolution..."
    
    # Copy downloaded packages to chroot
    cp "$BUILD_DIR/packages"/*.deb "$CHROOT_DIR/tmp/" || print_error "Failed to copy downloaded packages to chroot."
    
    # Install packages one by one using apt to handle dependencies properly
    print_status "Installing packages one by one with apt..."
    
    print_status "Installing libostree-1-1..."
    chroot "$CHROOT_DIR" apt install -y /tmp/libostree-1-1_2025.2-1~noble1_amd64.deb || print_error "Failed to install libostree-1-1"
    
    print_status "Installing ostree..."
    chroot "$CHROOT_DIR" apt install -y /tmp/ostree_2025.2-1~noble1_amd64.deb || print_error "Failed to install ostree"
    
    print_status "Installing bootc (dependencies already installed)..."
    chroot "$CHROOT_DIR" apt install -y /tmp/bootc_1.5.1-1~noble1_amd64.deb || print_error "Failed to install bootc"
    
    print_status "Installing apt-ostree..."
    chroot "$CHROOT_DIR" apt install -y /tmp/apt-ostree_0.1.0-1_amd64.deb || print_error "Failed to install apt-ostree"
    
    # Clean up
    chroot "$CHROOT_DIR" rm -f /tmp/*.deb || print_warning "Failed to clean up package files"
    
    # Verify installations with better debugging
    print_status "Verifying ostree and bootc installations..."
    
    # Debug: Check if ostree binary exists
    print_status "Checking ostree binary location..."
    chroot "$CHROOT_DIR" ls -la /usr/bin/ostree || print_warning "ostree binary not found at /usr/bin/ostree"
    chroot "$CHROOT_DIR" which ostree || print_warning "ostree not found in PATH"
    
    if chroot "$CHROOT_DIR" command -v ostree >/dev/null 2>&1; then
        print_success "ostree successfully installed and verified!"
    else
        print_warning "ostree command not found in PATH, but binary might exist"
        # Don't exit - the binary might be there but not in PATH
    fi
    
    if chroot "$CHROOT_DIR" command -v bootc >/dev/null 2>&1; then
        print_success "bootc successfully installed and verified!"
    else
        print_warning "bootc command not found in PATH, but binary might exist"
        # Don't exit - the binary might be there but not in PATH
    fi
    
    if chroot "$CHROOT_DIR" command -v apt-ostree >/dev/null 2>&1; then
        print_success "apt-ostree successfully installed and verified!"
    else
        print_warning "apt-ostree command not found in PATH, but binary might exist"
        # Don't exit - the binary might be there but not in PATH
    fi
    
    print_success "All critical ParticleOS packages installed (verification warnings above)"
else
    print_error "No downloaded packages found - CRITICAL: ostree, bootc, and apt-ostree are REQUIRED for ParticleOS!"
fi

# Enhanced snap removal and blocking
print_status "Removing snapd and setting APT preferences to block snaps..."
chroot "$CHROOT_DIR" bash -c 'DEBIAN_FRONTEND=noninteractive apt purge -y snapd ubuntu-advantage-tools update-notifier update-manager unattended-upgrades' || print_error "Failed to purge snapd and related packages."
chroot "$CHROOT_DIR" apt autoremove -y || print_error "Failed to autoremove orphaned packages."
chroot "$CHROOT_DIR" apt-mark hold snapd || print_error "Failed to hold snapd."

# Create an APT preference file to block snapd
chroot "$CHROOT_DIR" bash -c '
echo "Package: snapd" > /etc/apt/preferences.d/nosnap.pref && \
echo "Pin: release *" >> /etc/apt/preferences.d/nosnap.pref && \
echo "Pin-Priority: -1" >> /etc/apt/preferences.d/nosnap.pref
' || print_error "Failed to create APT preference file."

# Note: Firefox installation moved to host-side package download approach
print_status "Firefox will be installed from host-side downloaded packages"

# Enhanced system configuration
print_status "Configuring system settings..."
chroot "$CHROOT_DIR" bash -c '
echo "particleos" > /etc/hostname && \
echo "127.0.0.1 localhost" >> /etc/hosts && \
echo "127.0.1.1 particleos.local particleos" >> /etc/hosts && \
ln -sf /usr/share/zoneinfo/UTC /etc/localtime && \
locale-gen en_US.UTF-8 && \
update-locale LANG=en_US.UTF-8
' || print_error "Failed to configure system settings."

# Create user
print_status "Creating particle user..."
chroot "$CHROOT_DIR" useradd -m -s /bin/bash particle || print_error "Failed to create particle user."
# Skip password setting in chroot (PAM issues) - user will have NOPASSWD sudo access
print_warning "Skipping password setting for particle user (PAM issues in chroot). User will have NOPASSWD sudo access."
chroot "$CHROOT_DIR" usermod -aG sudo particle || print_error "Failed to add particle to sudo group."
chroot "$CHROOT_DIR" bash -c "echo 'particle ALL=(ALL) NOPASSWD:ALL' > /etc/sudoers.d/particle" || print_error "Failed to create sudoers file for particle."
chroot "$CHROOT_DIR" chmod 0440 /etc/sudoers.d/particle || print_error "Failed to set permissions for sudoers file."

# Enhanced service configuration
print_status "Enabling essential services..."
chroot "$CHROOT_DIR" systemctl enable sddm NetworkManager ssh || print_error "Failed to enable essential services."
chroot "$CHROOT_DIR" systemctl disable apt-daily.timer apt-daily-upgrade.timer || print_error "Failed to disable apt daily timers."

# Enhanced apt-ostree configuration
print_status "Configuring apt-ostree..."
chroot "$CHROOT_DIR" mkdir -p /etc/apt-ostree || print_error "Failed to create apt-ostree config directory."
chroot "$CHROOT_DIR" bash -c 'echo "ref: particleos/desktop/1.0.0" > /etc/apt-ostree/ref' || print_error "Failed to configure apt-ostree ref."

# Clean up
print_status "Cleaning APT caches and temporary files in chroot..."
chroot "$CHROOT_DIR" apt clean || print_error "Failed to clean apt cache in chroot."
chroot "$CHROOT_DIR" rm -rf /var/lib/apt/lists/* /tmp/* || print_error "Failed to remove apt lists or tmp files in chroot."

print_success "Base system configured"

print_header "Phase 3: Create ISO"

# Unmount filesystems from chroot before creating squashfs
cleanup_chroot_mounts

# Create ISO directory structure
mkdir -p "$ISO_DIR"/{casper,boot/grub,EFI/BOOT,isolinux} || print_error "Failed to create ISO directory structure."

# Copy kernel and initrd - more robust lookup
print_status "Copying kernel and initramfs..."
# Find the latest kernel and initrd
KERNEL_PATH=$(find "$CHROOT_DIR/boot" -maxdepth 1 -name "vmlinuz-*" | sort -V | tail -n 1)
INITRD_PATH=$(find "$CHROOT_DIR/boot" -maxdepth 1 -name "initrd.img-*" | sort -V | tail -n 1)

if [ -z "$KERNEL_PATH" ]; then
    print_error "Kernel (vmlinuz-*) not found in $CHROOT_DIR/boot."
fi
if [ -z "$INITRD_PATH" ]; then
    print_error "Initrd (initrd.img-*) not found in $CHROOT_DIR/boot."
fi

cp "$KERNEL_PATH" "$ISO_DIR/casper/vmlinuz" || print_error "Failed to copy kernel."
cp "$INITRD_PATH" "$ISO_DIR/casper/initrd" || print_error "Failed to copy initrd."
print_success "Kernel and initramfs copied."

# Create filesystem manifest
print_status "Creating filesystem manifest..."
chroot "$CHROOT_DIR" dpkg-query -W --showformat='${Package} ${Version}\n' > "/filesystem.manifest.tmp" || print_error "Failed to create filesystem manifest in chroot."
mv "/filesystem.manifest.tmp" "$ISO_DIR/casper/filesystem.manifest" || print_error "Failed to move filesystem manifest."
print_success "Filesystem manifest created."

# Create filesystem size
print_status "Calculating filesystem size..."
du -sx --block-size=1 "$CHROOT_DIR" | cut -f1 > "$ISO_DIR/casper/filesystem.size" || print_error "Failed to create filesystem.size."
print_success "Filesystem size calculated."

# Create squashfs
print_status "Creating squashfs filesystem..."
# Ensure chroot mounts are truly clean before this step
if mountpoint -q "$CHROOT_DIR/dev" || mountpoint -q "$CHROOT_DIR/run" || mountpoint -q "$CHROOT_DIR/proc" || mountpoint -q "$CHROOT_DIR/sys"; then
    print_error "Chroot mounts are still active before mksquashfs. This should not happen."
fi
mksquashfs "$CHROOT_DIR" "$ISO_DIR/casper/filesystem.squashfs" -comp xz -e boot || print_error "Failed to create squashfs filesystem."
print_success "Squashfs created."

# Create GRUB configuration
cat > "$ISO_DIR/boot/grub/grub.cfg" << 'GRUBEOF'
set timeout=10
set default=0

menuentry "Try ParticleOS without installing" {
    linux /casper/vmlinuz boot=casper quiet splash ---
    initrd /casper/initrd
}

menuentry "Install ParticleOS" {
    linux /casper/vmlinuz boot=casper quiet splash ---
    initrd /casper/initrd
}

menuentry "Check disc for defects" {
    linux /casper/vmlinuz boot=casper integrity-check quiet splash ---
    initrd /casper/initrd
}
GRUBEOF
print_success "GRUB configuration created."


# Create ISOLINUX configuration
cat > "$ISO_DIR/isolinux/isolinux.cfg" << 'ISOLINUXEOF'
DEFAULT live
TIMEOUT 300
PROMPT 1

LABEL live
  MENU LABEL Try ParticleOS without installing
  KERNEL /casper/vmlinuz
  APPEND boot=casper initrd=/casper/initrd quiet splash ---

LABEL live-install
  MENU LABEL Install ParticleOS
  KERNEL /casper/vmlinuz
  APPEND boot=casper initrd=/casper/initrd quiet splash ---

LABEL check
  MENU LABEL Check disc for defects
  KERNEL /casper/vmlinuz
  APPEND boot=casper integrity-check initrd=/casper/initrd quiet splash ---
ISOLINUXEOF
print_success "ISOLINUX configuration created."

# Copy ISOLINUX boot files - corrected paths
print_status "Copying ISOLINUX boot files..."
ISOLINUX_BIN_HOST_PATH="/usr/lib/ISOLINUX/isolinux.bin"
ISOHDPFX_BIN_HOST_PATH="/usr/lib/ISOLINUX/isohdpfx.bin"

if [ ! -f "$ISOLINUX_BIN_HOST_PATH" ]; then print_error "isolinux.bin not found at $ISOLINUX_BIN_HOST_PATH. Install isolinux."; fi
if [ ! -f "$ISOHDPFX_BIN_HOST_PATH" ]; then print_error "isohdpfx.bin not found at $ISOHDPFX_BIN_HOST_PATH. Install isolinux."; fi

cp "$ISOLINUX_BIN_HOST_PATH" "$ISO_DIR/isolinux/" || print_error "Failed to copy isolinux.bin."
# Note: boot.cat is generated by xorriso during ISO creation, not copied from host
print_success "ISOLINUX boot files copied."

# Create EFI boot image (for UEFI) using grub-mkimage
print_status "Creating EFI boot image..."
GRUB_EFI_MODULES_DIR="/usr/lib/grub/x86_64-efi"
if [ ! -d "$GRUB_EFI_MODULES_DIR" ]; then
    print_error "GRUB EFI modules directory not found at $GRUB_EFI_MODULES_DIR. Cannot create EFI boot image. Check grub-efi-amd64-bin installation."
fi

# Ensure EFI/BOOT directory exists
mkdir -p "$ISO_DIR/EFI/BOOT" || print_error "Failed to create EFI/BOOT directory"

# Create EFI boot image directly
grub-mkimage \
    -o "$ISO_DIR/EFI/BOOT/bootx64.efi" \
    -p "/boot/grub" \
    -O x86_64-efi \
    boot linux normal configfile part_gpt part_msdos fat \
    squash4 test configfile search loadenv \
    efi_gop efi_uga all_video gfxterm_menu \
    chain iso9660 || print_error "Failed to create EFI boot image."

# Create a dummy efi.img file for xorriso compatibility
dd if=/dev/zero of="$ISO_DIR/boot/grub/efi.img" bs=1M count=1 2>/dev/null || print_error "Failed to create dummy efi.img"

print_success "EFI boot image created."


# Create ISO
print_status "Creating bootable ISO using xorriso..."
xorriso -as mkisofs \
    -o "$OUTPUT_DIR/${PROJECT_NAME}-${VERSION}.iso" \
    -J -joliet-long \
    -r -V "ParticleOS ${VERSION}" \
    -b isolinux/isolinux.bin \
    -c isolinux/boot.cat \
    -boot-load-size 4 -boot-info-table \
    -no-emul-boot -eltorito-alt-boot \
    -e boot/grub/efi.img -no-emul-boot \
    -isohybrid-mbr "$ISOHDPFX_BIN_HOST_PATH" \
    -partition_offset 16 \
    -part_like_isohybrid \
    "$ISO_DIR" || print_error "Failed to create ISO with xorriso."

print_success "ISO created successfully: $OUTPUT_DIR/${PROJECT_NAME}-${VERSION}.iso"

# Clean up temporary directories within the container
print_status "Cleaning up temporary directories in container..."
rm -rf "$CHROOT_DIR" "$ISO_DIR" || print_error "Failed to remove temporary directories in container."

print_header "Container Build Complete!"
echo "🎉 ParticleOS ISO built successfully in container!"
echo "📁 Location: $OUTPUT_DIR/${PROJECT_NAME}-${VERSION}.iso"
EOF
    
    chmod +x "$BUILD_DIR/build-in-container.sh"
    print_success "Container build script created"
}

# Build the container image
build_container_image() {
    print_header "Phase 5: Build Container Image"
    
    print_status "Building container image '$IMAGE_NAME'..."
    cd "$BUILD_DIR"
    
    # Pass apt-cacher-ng configuration as build argument if available
    local build_args=""
    if [ "${APT_CACHER_AVAILABLE:-false}" = "true" ]; then
        build_args="--build-arg APT_CACHER_URL=$APT_CACHER_URL"
        print_status "Building with apt-cacher-ng: $APT_CACHER_URL"
    else
        print_status "Building without cache (direct repository access)"
    fi
    
    # Use || print_error instead of full if/else for cleaner code
    $PODMAN_CMD build $build_args -t "$IMAGE_NAME" . || print_error "Failed to build container image '$IMAGE_NAME'."
    
    print_success "Container image built successfully"
}

# Run the build in the container
run_container_build() {
    print_header "Phase 6: Run Build in Container"
    
    print_status "Starting build in isolated container '$CONTAINER_NAME'..."
    
    # Pass apt-cacher-ng configuration to container if available
    local cache_env=""
    if [ "${APT_CACHER_AVAILABLE:-false}" = "true" ]; then
        cache_env="-e APT_CACHER_URL=$APT_CACHER_URL"
        print_status "Using apt-cacher-ng: $APT_CACHER_URL"
    else
        print_status "Using direct repository access (no cache)"
    fi
    
    # Run the container with volume mounts
    # Added --privileged for mmdebstrap/chroot/mount/mksquashfs within container.
    # --userns=host for better compatibility with privileged operations in some cases.
    $PODMAN_CMD run \
        --name "$CONTAINER_NAME" \
        --rm \
        --privileged \
        --userns=host \
        $cache_env \
        -v "$OUTPUT_DIR:/output:Z" \
        -v "$SCRIPT_DIR/pkgs:/host/pkgs:Z" \
        "$IMAGE_NAME" || print_error "Container build failed. Check container logs for details."
    
    print_success "Container build completed successfully"
}

# Cleanup function (for host-side Podman resources)
cleanup() {
    print_status "Cleaning up Podman container and image..."
    
    # Remove container if it exists
    if $PODMAN_CMD container exists "$CONTAINER_NAME" 2>/dev/null; then
        print_status "Removing container: $CONTAINER_NAME"
        $PODMAN_CMD container rm -f "$CONTAINER_NAME" 2>/dev/null || true
    fi
    
    # Remove image if it exists. Be careful here, sometimes you want to keep the image.
    # For a clean build process, removing it ensures fresh start next time.
    if $PODMAN_CMD image exists "$IMAGE_NAME" 2>/dev/null; then
        print_status "Removing image: $IMAGE_NAME"
        $PODMAN_CMD image rm -f "$IMAGE_NAME" 2>/dev/null || true
    fi
    
    print_success "Podman resources cleaned."
}

# Signal trap for cleanup
trap cleanup EXIT INT TERM

# Main execution
main() {
    echo "🚀 ParticleOS ISO Builder (Podman) - HARDENED ULTRA SAFE VERSION"
    echo "=============================================================="
    echo "Project: $PROJECT_NAME"
    echo "Version: $VERSION"
    echo "Build Directory: $BUILD_DIR"
    echo "Tool: Podman (Containerized)"
    echo ""
    echo "🛡️  This build runs in a completely isolated container"
    echo "🛡️  Your host system cannot be affected by this build"
    echo "🛡️  Includes all hardening features from the analysis"
    echo "🛡️  Firefox installed as DEB package (no snap)"
    echo "🚀 apt-cacher-ng support for faster builds"
    echo ""

    # Call cleanup explicitly at the start to ensure a clean slate before attempting new build
    cleanup # This will also remove image from previous failed runs.

    check_podman
    check_prerequisites
    clean_build # This cleans host-side directories, but cleanup() handles Podman resources.
    create_dockerfile
    create_container_build_script
    build_container_image
    run_container_build

    print_header "Build Complete!"
    echo ""
    echo "🎉 ParticleOS ISO built successfully!"
    echo "📁 Location: $OUTPUT_DIR/${PROJECT_NAME}-${VERSION}.iso"
    echo ""
    echo "🧪 Test the ISO:"
    echo "    qemu-system-x86_64 -m 4G -enable-kvm \\"
    echo "      -cdrom $OUTPUT_DIR/${PROJECT_NAME}-${VERSION}.iso \\"
    echo "      -boot d"
    echo ""
    echo "🛡️  Your host system is completely safe!"
}

# Run main function
main "$@"