#!/bin/bash

# deb-bootc-image-builder - Simplified container to disk image converter
# This is a simplified fork of the original Fedora bootc-image-builder
# that uses native Debian tools instead of osbuild

set -euo pipefail

# System tool paths
QEMU_IMG="/usr/bin/qemu-img"
RSYNC="/usr/bin/rsync"
TAR="/bin/tar"
MKDIR="/bin/mkdir"
CP="/bin/cp"
LOSETUP="/usr/sbin/losetup"
PARTED="/usr/sbin/parted"
MOUNT="/bin/mount"
UMOUNT="/bin/umount"
BOOTUPD="/usr/libexec/bootupd"
GRUB_INSTALL="/usr/sbin/grub-install"
GRUB_MKCONFIG="/usr/sbin/grub-mkconfig"

# Default values
CONTAINER_IMAGE=""
OUTPUT_FORMAT="qcow2"
OUTPUT_DIR="/output"
STORE_DIR="/store"
CONTAINER_STORAGE="/var/lib/containers/storage"

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

# Global variables for cleanup
WORK_DIR=""
LOOP_DEV=""
MOUNT_POINT=""
CONTAINER_FS_DIR=""
DISK_RAW=""

# Logging functions
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Help function
show_help() {
    cat << EOF
Simplified Debian bootc-image-builder

Usage: $0 [OPTIONS] <container-image>

Options:
    -f, --format FORMAT    Output format (qcow2, raw, img) [default: qcow2]
    -o, --output DIR       Output directory [default: /output]
    -s, --store DIR        Store directory [default: /store]
    -h, --help            Show this help message

Examples:
    $0 simple-cli:latest
    $0 -f raw -o /tmp/output simple-cli:latest
    $0 --format qcow2 --output /output simple-cli:latest

EOF
}

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        -f|--format)
            OUTPUT_FORMAT="$2"
            shift 2
            ;;
        -o|--output)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        -s|--store)
            STORE_DIR="$2"
            shift 2
            ;;
        -h|--help)
            show_help
            exit 0
            ;;
        -*)
            log_error "Unknown option $1"
            show_help
            exit 1
            ;;
        *)
            CONTAINER_IMAGE="$1"
            shift
            ;;
    esac
done

# Validate inputs
if [[ -z "$CONTAINER_IMAGE" ]]; then
    log_error "Container image is required"
    show_help
    exit 1
fi

if [[ ! -d "$OUTPUT_DIR" ]]; then
    log_error "Output directory $OUTPUT_DIR does not exist"
    exit 1
fi

# Validate output format
case "$OUTPUT_FORMAT" in
    qcow2|raw|img)
        ;;
    *)
        log_error "Unsupported output format: $OUTPUT_FORMAT"
        exit 1
        ;;
esac

log_info "Starting simplified bootc-image-builder"
log_info "Container image: $CONTAINER_IMAGE"
log_info "Output format: $OUTPUT_FORMAT"
log_info "Output directory: $OUTPUT_DIR"

# Create temporary working directory
# Use /var/tmp or /home instead of /tmp to avoid tmpfs space limitations
WORK_DIR=$(mktemp -d /var/tmp/bootc-builder-XXXXXX 2>/dev/null || mktemp -d /home/joe/tmp/bootc-builder-XXXXXX 2>/dev/null || echo "/home/joe/tmp/bootc-builder-$$")
if [[ ! -d "$WORK_DIR" ]]; then
    mkdir -p "$WORK_DIR" || {
        log_error "Failed to create working directory"
        exit 1
    }
fi
log_info "Working directory: $WORK_DIR"

# Check available space
AVAILABLE_SPACE=$(df "$WORK_DIR" | awk 'NR==2 {print $4}')
if [[ "$AVAILABLE_SPACE" -lt 5000000 ]]; then  # Less than 5GB
    log_warn "Low disk space available: ${AVAILABLE_SPACE}KB"
    log_warn "This might cause issues with large containers"
fi

# Signal handling for cleanup
cleanup_on_exit() {
    local exit_code=$?
    log_info "Cleaning up resources..."
    
    # Unmount if mounted
    if [[ -n "$MOUNT_POINT" && -d "$MOUNT_POINT" ]]; then
        if mountpoint -q "$MOUNT_POINT" 2>/dev/null; then
            log_info "Unmounting $MOUNT_POINT..."
            sudo umount "$MOUNT_POINT" 2>/dev/null || true
        fi
    fi
    
    # Detach loop device if exists
    if [[ -n "$LOOP_DEV" && -b "$LOOP_DEV" ]]; then
        log_info "Detaching loop device $LOOP_DEV..."
        sudo losetup -d "$LOOP_DEV" 2>/dev/null || true
    fi
    
    # Clean up any remaining loop devices that might be ours
    if [[ -n "$WORK_DIR" ]]; then
        log_info "Checking for orphaned loop devices..."
        for loop_dev in /dev/loop*; do
            if [[ -b "$loop_dev" ]]; then
                local mount_point
                mount_point=$(losetup -l "$loop_dev" 2>/dev/null | grep -o "/var/tmp/bootc-builder-[^[:space:]]*" || true)
                if [[ -n "$mount_point" ]]; then
                    log_info "Detaching orphaned loop device $loop_dev..."
                    sudo losetup -d "$loop_dev" 2>/dev/null || true
                fi
            fi
        done
    fi
    
    # Remove working directory
    if [[ -n "$WORK_DIR" && -d "$WORK_DIR" ]]; then
        log_info "Removing working directory $WORK_DIR..."
        sudo rm -rf "$WORK_DIR" 2>/dev/null || true
    fi
    
    # Remove any remaining temporary files
    if [[ -n "$DISK_RAW" && -f "$DISK_RAW" ]]; then
        log_info "Removing temporary disk image $DISK_RAW..."
        sudo rm -f "$DISK_RAW" 2>/dev/null || true
    fi
    
    if [[ -n "$CONTAINER_FS_DIR" && -d "$CONTAINER_FS_DIR" ]]; then
        log_info "Removing container filesystem directory $CONTAINER_FS_DIR..."
        sudo rm -rf "$CONTAINER_FS_DIR" 2>/dev/null || true
    fi
    
    log_info "Cleanup completed"
    exit "$exit_code"
}

# Set up signal handlers
trap cleanup_on_exit EXIT INT TERM

# Function to clean up any existing loop devices from previous runs
cleanup_existing_loop_devices() {
    log_info "Checking for existing loop devices from previous runs..."
    for loop_dev in /dev/loop*; do
        if [[ -b "$loop_dev" ]]; then
            local mount_point
            mount_point=$(losetup -l "$loop_dev" 2>/dev/null | grep -o "/var/tmp/bootc-builder-[^[:space:]]*" || true)
            if [[ -n "$mount_point" ]]; then
                log_warn "Found existing loop device $loop_dev from previous run, cleaning up..."
                # Try to unmount if mounted
                if mountpoint -q "$mount_point" 2>/dev/null; then
                    sudo umount "$mount_point" 2>/dev/null || true
                fi
                # Detach the loop device
                sudo losetup -d "$loop_dev" 2>/dev/null || true
                # Remove the mount point directory
                sudo rm -rf "$mount_point" 2>/dev/null || true
            fi
        fi
    done
}

# Clean up any existing loop devices at startup
cleanup_existing_loop_devices

# Check if container image exists locally
log_info "Checking if container image exists locally: $CONTAINER_IMAGE"
if podman images --format "{{.Repository}}:{{.Tag}}" | grep -q "^$CONTAINER_IMAGE$"; then
    log_info "Container image found locally, skipping pull"
elif podman images --format "{{.Repository}}:{{.Tag}}" | grep -q "simple-cli:latest"; then
    log_info "Found simple-cli:latest locally, using that instead"
    CONTAINER_IMAGE="simple-cli:latest"
elif podman images --format "{{.Repository}}:{{.Tag}}" | grep -q "localhost/simple-cli:latest"; then
    log_info "Found localhost/simple-cli:latest locally, using that instead"
    CONTAINER_IMAGE="localhost/simple-cli:latest"
else
    log_info "Container image not found locally, attempting to pull: $CONTAINER_IMAGE"
    podman pull "$CONTAINER_IMAGE" || {
        log_error "Failed to pull container image. Please ensure the image is available."
        log_error "You can either:"
        log_error "1. Build the image locally first"
        log_error "2. Use a fully qualified registry URL"
        exit 1
    }
fi

# Export the container to a tar file
log_info "Exporting container to tar..."
podman export "$(podman create --rm "$CONTAINER_IMAGE" true)" > "$WORK_DIR/container.tar"

# Create a loopback device and mount the container
log_info "Setting up loopback device..."
cd "$WORK_DIR"

# Extract the container
log_info "Extracting container filesystem..."
CONTAINER_FS_DIR="$WORK_DIR/container-fs"
mkdir -p "$CONTAINER_FS_DIR"
if ! tar -xf container.tar -C "$CONTAINER_FS_DIR"; then
    log_error "Failed to extract container. This might be due to:"
    log_error "1. Insufficient disk space"
    log_error "2. Corrupted container archive"
    log_error "3. Permission issues"
    exit 1
fi

# Create disk image using qemu-img (no sudo required)
log_info "Creating disk image..."
IMAGE_SIZE="2G"  # Default size, could be made configurable

# Create a proper disk image with partitions for bootloader installation
log_info "Creating partitioned disk image for bootloader installation..."

# Create raw image file
DISK_RAW="$WORK_DIR/disk.raw"
$QEMU_IMG create -f raw "$DISK_RAW" "$IMAGE_SIZE"

# For bootloader installation, we need to create a proper disk structure
# This requires sudo for partitioning and mounting
log_info "Setting up disk partitions for bootloader installation..."
log_info "Note: This step requires sudo for disk operations"

# Create partition table and filesystem
log_info "Creating partition table and filesystem..."
sudo $PARTED "$DISK_RAW" mklabel msdos
sudo $PARTED "$DISK_RAW" mkpart primary ext4 1MiB 100%
sudo $PARTED "$DISK_RAW" set 1 boot on

# Setup loopback device
log_info "Setting up loopback device..."
LOOP_DEV=$(sudo $LOSETUP --find --show "$DISK_RAW")
log_info "Using loopback device: $LOOP_DEV"

# Force kernel to reread partition table
log_info "Rereading partition table..."
sudo partprobe "$LOOP_DEV" || true

# Get the partition device
PART_DEV="${LOOP_DEV}p1"
log_info "Partition device: $PART_DEV"

# Wait for partition to be available and verify it exists
log_info "Waiting for partition to be available..."
for i in {1..10}; do
    if [[ -b "$PART_DEV" ]]; then
        log_info "Partition device found: $PART_DEV"
        break
    fi
    log_info "Waiting for partition device... (attempt $i/10)"
    sleep 1
done

if [[ ! -b "$PART_DEV" ]]; then
    log_error "Partition device not found: $PART_DEV"
    log_error "Available devices:"
    ls -la /dev/loop* || true
    exit 1
fi

# Create filesystem
log_info "Creating ext4 filesystem..."
sudo mkfs.ext4 "$PART_DEV"

# Mount the partition
MOUNT_POINT="$WORK_DIR/mount-point"
mkdir -p "$MOUNT_POINT"
sudo $MOUNT "$PART_DEV" "$MOUNT_POINT"

# Copy container filesystem
log_info "Copying container filesystem to disk image..."
sudo rsync -a "$CONTAINER_FS_DIR/" "$MOUNT_POINT/"

# Install bootloader manually since deb-bootupd doesn't support non-default sysroots
log_info "Installing bootloader manually..."
if [[ -x "$GRUB_INSTALL" ]]; then
    log_info "Found grub-install, installing GRUB bootloader..."
    
    # Install GRUB to the mounted filesystem
    # Use --root-directory to specify the mount point
    # Use --target=i386-pc for BIOS systems
    # Use --boot-directory to specify where boot files are located
    log_info "Installing GRUB to $MOUNT_POINT..."
    sudo "$GRUB_INSTALL" \
        --root-directory="$MOUNT_POINT" \
        --target=i386-pc \
        --boot-directory="$MOUNT_POINT/boot" \
        --force \
        "$LOOP_DEV" || {
        log_warn "GRUB installation failed, continuing without bootloader"
    }
    
    # Generate GRUB configuration manually since grub-mkconfig needs host environment
    log_info "Generating GRUB configuration manually..."
    
    # Create a basic GRUB configuration for the container
    cat > "$MOUNT_POINT/boot/grub/grub.cfg" << 'EOF'
# GRUB configuration for simple-cli container
set timeout=5
set default=0

menuentry "Simple CLI" {
    set root=(hd0,msdos1)
    linux /boot/vmlinuz-6.12.38+deb13-amd64 root=/dev/sda1 rw console=ttyS0
    initrd /boot/initrd.img-6.12.38+deb13-amd64
}

menuentry "Simple CLI (Recovery)" {
    set root=(hd0,msdos1)
    linux /boot/vmlinuz-6.12.38+deb13-amd64 root=/dev/sda1 rw single console=ttyS0
    initrd /boot/initrd.img-6.12.38+deb13-amd64
}
EOF
    
    log_info "GRUB configuration created at $MOUNT_POINT/boot/grub/grub.cfg"
    
else
    log_warn "grub-install not found at $GRUB_INSTALL, skipping bootloader installation"
    log_warn "Image will not be bootable without bootloader"
fi

# Unmount
sudo $UMOUNT "$MOUNT_POINT"

# Detach loopback device
sudo $LOSETUP -d "$LOOP_DEV"

# Convert to requested format
log_info "Converting to $OUTPUT_FORMAT format..."
case "$OUTPUT_FORMAT" in
    qcow2)
        $QEMU_IMG convert -f raw -O qcow2 "$DISK_RAW" "$OUTPUT_DIR/$(basename "$CONTAINER_IMAGE" | tr ':' '_').qcow2"
        ;;
    raw)
        cp "$DISK_RAW" "$OUTPUT_DIR/$(basename "$CONTAINER_IMAGE" | tr ':' '_').raw"
        ;;
    img)
        cp "$DISK_RAW" "$OUTPUT_DIR/$(basename "$CONTAINER_IMAGE" | tr ':' '_').img"
        ;;
esac

log_info "Image creation completed successfully!"
log_info "Output file: $OUTPUT_DIR/$(basename "$CONTAINER_IMAGE" | tr ':' '_').$OUTPUT_FORMAT"

# List output files
log_info "Output directory contents:"
ls -la "$OUTPUT_DIR"