#!/bin/bash

# ParticleOS ISO Builder with mmdebstrap
# Builds a bootable ISO using mmdebstrap + apt-ostree integration

set -e

# 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"
}

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"
CHROOT_DIR="$BUILD_DIR/chroot"
ISO_DIR="$BUILD_DIR/iso"
OUTPUT_DIR="$SCRIPT_DIR/output"

# Check prerequisites
check_prerequisites() {
    print_header "Phase 1: Check Prerequisites"
    
    local missing_packages=()
    
    for package in mmdebstrap squashfs-tools xorriso grub-pc-bin grub-efi-amd64-bin; do
        if ! dpkg -l | grep -q "^ii  $package "; then
            missing_packages+=("$package")
        fi
    done
    
    if [ ${#missing_packages[@]} -gt 0 ]; then
        print_status "Installing missing packages: ${missing_packages[*]}"
        sudo apt update
        sudo apt install -y "${missing_packages[@]}"
    fi
    
    print_success "All prerequisites satisfied"
}

# Clean build environment
clean_build() {
    print_header "Phase 2: Clean Build Environment"
    
    if [ -d "$BUILD_DIR" ]; then
        print_status "Removing previous build directory..."
        sudo rm -rf "$BUILD_DIR"
    fi
    
    mkdir -p "$BUILD_DIR" "$CHROOT_DIR" "$ISO_DIR" "$OUTPUT_DIR"
    print_success "Build environment cleaned"
}

# Create base system using mmdebstrap
create_base_system() {
    print_header "Phase 3: Create Base System with mmdebstrap"
    
    print_status "Creating base Ubuntu system using mmdebstrap..."
    
    # Create base system with mmdebstrap
    sudo mmdebstrap \
        --architectures=amd64 \
        --variant=minbase \
        --include=systemd,systemd-sysv,dbus,curl,ca-certificates \
        noble \
        "$CHROOT_DIR" \
        http://archive.ubuntu.com/ubuntu/
    
    if [ $? -eq 0 ]; then
        print_success "Base system created with mmdebstrap"
    else
        print_error "Failed to create base system"
        exit 1
    fi
}

# Configure base system
configure_base_system() {
    print_header "Phase 4: Configure Base System"
    
    print_status "Configuring base system..."
    
    # Mount necessary filesystems
    sudo mount --bind /dev "$CHROOT_DIR/dev"
    sudo mount --bind /run "$CHROOT_DIR/run"
    sudo mount -t proc none "$CHROOT_DIR/proc"
    sudo mount -t sysfs none "$CHROOT_DIR/sys"
    
    # Configure package sources
    sudo chroot "$CHROOT_DIR" bash -c "echo 'deb http://archive.ubuntu.com/ubuntu noble main restricted universe multiverse' > /etc/apt/sources.list"
    sudo chroot "$CHROOT_DIR" bash -c "echo 'deb http://archive.ubuntu.com/ubuntu noble-updates main restricted universe multiverse' >> /etc/apt/sources.list"
    sudo chroot "$CHROOT_DIR" bash -c "echo 'deb http://security.ubuntu.com/ubuntu noble-security main restricted universe multiverse' >> /etc/apt/sources.list"
    
    # Add your Forgejo repository for apt-ostree
    sudo chroot "$CHROOT_DIR" bash -c "curl -o /etc/apt/keyrings/forgejo-robojerk.asc https://git.raines.xyz/api/packages/robojerk/debian/gpg"
    sudo chroot "$CHROOT_DIR" bash -c "echo 'deb [signed-by=/etc/apt/keyrings/forgejo-robojerk.asc] https://git.raines.xyz/api/packages/robojerk/debian noble main' > /etc/apt/sources.list.d/forgejo.list"
    
    # Update package lists
    sudo chroot "$CHROOT_DIR" apt update
    
    # Install desktop and additional packages
    sudo chroot "$CHROOT_DIR" apt install -y \
        kubuntu-desktop \
        plasma-desktop plasma-workspace kde-plasma-desktop sddm \
        ostree bootc \
        flatpak \
        network-manager plasma-nm \
        openssh-server \
        curl wget vim nano \
        htop fastfetch tree \
        firefox \ # Firefox is a snap package, so we need to install it from the custom repository
        pulseaudio pulseaudio-utils \
        fonts-ubuntu fonts-noto \
        build-essential git
    
    # Download and install apt-ostree from custom repository
    print_status "Installing apt-ostree from custom repository..."
    sudo chroot "$CHROOT_DIR" timeout 60 wget -O /tmp/apt-ostree.deb "https://git.raines.xyz/robojerk/apt-ostree/raw/branch/main/apt-ostree_0.1.0-1_amd64.deb"
    sudo chroot "$CHROOT_DIR" dpkg -i /tmp/apt-ostree.deb
    sudo chroot "$CHROOT_DIR" rm /tmp/apt-ostree.deb
    
    # Remove unwanted packages and block snaps
    sudo chroot "$CHROOT_DIR" apt remove -y snapd ubuntu-advantage-tools update-notifier update-manager unattended-upgrades
    sudo chroot "$CHROOT_DIR" apt-mark hold snapd
    sudo chroot "$CHROOT_DIR" bash -c "echo 'Package: snapd' > /etc/apt/preferences.d/no-snapd"
    sudo chroot "$CHROOT_DIR" bash -c "echo 'Pin: release *' >> /etc/apt/preferences.d/no-snapd"
    sudo chroot "$CHROOT_DIR" bash -c "echo 'Pin-Priority: -1' >> /etc/apt/preferences.d/no-snapd"
    
    # Configure system
    sudo chroot "$CHROOT_DIR" bash -c "echo 'particleos' > /etc/hostname"
    sudo chroot "$CHROOT_DIR" bash -c "echo '127.0.1.1 particleos' >> /etc/hosts"
    
    # Create user
    sudo chroot "$CHROOT_DIR" useradd -m -s /bin/bash -G sudo particle
    sudo chroot "$CHROOT_DIR" bash -c "echo 'particle:particle' | chpasswd"
    
    # Enable services
    sudo chroot "$CHROOT_DIR" systemctl enable sddm
    sudo chroot "$CHROOT_DIR" systemctl enable NetworkManager
    sudo chroot "$CHROOT_DIR" systemctl enable ssh
    
    # Configure apt-ostree
    sudo chroot "$CHROOT_DIR" bash -c "mkdir -p /etc/apt-ostree"
    sudo chroot "$CHROOT_DIR" bash -c "echo 'ref: particleos/desktop/1.0.0' > /etc/apt-ostree/ref"
    
    # Unmount filesystems
    sudo umount "$CHROOT_DIR/dev"
    sudo umount "$CHROOT_DIR/run"
    sudo umount "$CHROOT_DIR/proc"
    sudo umount "$CHROOT_DIR/sys"
    
    print_success "Base system configured"
}

# Create live filesystem
create_live_fs() {
    print_header "Phase 5: Create Live Filesystem"
    
    print_status "Creating live filesystem..."
    
    # Create ISO directory structure
    mkdir -p "$ISO_DIR/casper"
    mkdir -p "$ISO_DIR/boot/grub"
    mkdir -p "$ISO_DIR/isolinux"
    
    # Create squashfs from the chroot
    sudo mksquashfs "$CHROOT_DIR" "$ISO_DIR/casper/filesystem.squashfs" -comp xz -e boot
    
    # Create filesystem.manifest
    sudo chroot "$CHROOT_DIR" dpkg-query -W --showformat='${Package} ${Version}\n' > "$ISO_DIR/casper/filesystem.manifest"
    
    # Create filesystem.size
    sudo du -sx --block-size=1 "$CHROOT_DIR" | cut -f1 > "$ISO_DIR/casper/filesystem.size"
    
    # Copy kernel and initramfs
    sudo cp "$CHROOT_DIR/boot/vmlinuz-"* "$ISO_DIR/casper/vmlinuz"
    sudo cp "$CHROOT_DIR/boot/initrd.img-"* "$ISO_DIR/casper/initrd"
    
    print_success "Live filesystem created"
}

# Setup boot configuration
setup_boot() {
    print_header "Phase 6: Setup Boot Configuration"
    
    print_status "Setting up boot configuration..."
    
    # Create GRUB configuration
    cat > "$ISO_DIR/boot/grub/grub.cfg" << 'EOF'
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
}
EOF
    
    # Create ISOLINUX configuration
    cat > "$ISO_DIR/isolinux/isolinux.cfg" << 'EOF'
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 --
EOF
    
    print_success "Boot configuration setup complete"
}

# Create ISO
create_iso() {
    print_header "Phase 7: Create ISO"
    
    print_status "Creating bootable ISO..."
    
    # Create ISO using xorriso
    xorriso -as mkisofs \
        -o "$OUTPUT_DIR/${PROJECT_NAME}-${VERSION}.iso" \
        -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 /usr/lib/ISOLINUX/isohdpfx.bin \
        -r -V "ParticleOS ${VERSION}" \
        "$ISO_DIR"
    
    if [ $? -eq 0 ]; then
        print_success "ISO created successfully: $OUTPUT_DIR/${PROJECT_NAME}-${VERSION}.iso"
    else
        print_error "Failed to create ISO"
        exit 1
    fi
}

# Main build process
main() {
    echo "🚀 ParticleOS ISO Builder (mmdebstrap)"
    echo "======================================="
    echo "Project: $PROJECT_NAME"
    echo "Version: $VERSION"
    echo "Build Directory: $BUILD_DIR"
    echo "Tool: mmdebstrap"
    echo ""
    
    # Run build phases
    check_prerequisites
    clean_build
    create_base_system
    configure_base_system
    create_live_fs
    setup_boot
    create_iso
    
    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 ""
}

# Run main function
main "$@"